Trace sparse matrices in linear functions (part 4)

tum-pbs · Nov 10, 2023 · c87517a · c87517a
1 parent e1fdadf
commit c87517a
Show file tree

Hide file tree

Showing 2 changed files with 253 additions and 119 deletions.
diff --git a/phiml/math/_ops.py b/phiml/math/_ops.py
@@ -2314,7 +2314,7 @@ def gather(values: Tensor, indices: Tensor, dims: Union[DimFilter, None] = None)
             indices = expand(indices, channel(gather=dims))
         if not channel(indices).item_names[0]:
             indices = indices._with_shape_replaced(indices.shape.with_dim_size(channel(indices), dims))
-        return values.to_sparse_tracer().gather(indices)
+        return values._gather(indices)
     treat_as_batch = non_channel(indices).only(values.shape).without(dims)
     batch_ = (values.shape.batch & indices.shape.batch).without(dims) & treat_as_batch
     channel_ = values.shape.without(dims).without(batch_)