Fix argmax,argmin dim checks

tum-pbs · Aug 9, 2024 · 9858107 · 9858107
1 parent 79322b6
commit 9858107
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/phiml/math/_ops.py b/phiml/math/_ops.py
@@ -1700,7 +1700,7 @@ def argmax(x: Tensor, dim: DimFilter, index_dim=channel('index')):
     """
     dims = x.shape.only(dim)
     keep = x.shape.without(dims)
-    assert dim, f"No dimensions {dim} present on key {x.shape}"
+    assert dims, f"argmax requires dim to be present on data but {dim} does not exist on {x.shape}"
     if isinstance(x, (SparseCoordinateTensor, CompressedSparseMatrix)):
         if dims in sparse_dims(x):
             max_val = max_(x, dim)
@@ -1740,7 +1740,7 @@ def argmin(x: Tensor, dim: DimFilter, index_dim=channel('index')):
     """
     dims = x.shape.only(dim)
     keep = x.shape.without(dims)
-    assert dim, f"No dimensions {dim} present on key {x.shape}"
+    assert dims, f"argmin requires dim to be present on data but {dim} does not exist on {x.shape}"
     if isinstance(x, (SparseCoordinateTensor, CompressedSparseMatrix)):
         if dims in sparse_dims(x):
             min_val = min_(x, dim)