Support indices in Shape.only()

tum-pbs · Feb 8, 2024 · 2249df6 · 2249df6
1 parent 741e7fc
commit 2249df6
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/phiml/math/_shape.py b/phiml/math/_shape.py
@@ -779,6 +779,10 @@ def only(self, dims: 'DimFilter', reorder=False):
         if isinstance(dims, Shape):
             dims = dims.names
         if isinstance(dims, (tuple, list, set)):
+            if all(isinstance(d, int) for d in dims):
+                if not reorder:
+                    dims = tuple(sorted(dims))
+                return self[dims]
             dim_names = []
             for d in dims:
                 if callable(d):