mindtorch_v2: align all/any/count_nonzero dim schema errors (#2751)

Author: lvyufeng

src/mindtorch_v2/_dispatch/schema.py

@@ -292,6 +292,64 @@ def _validate_arg_reduce_dim(value):
                     f"{op_short_name}(): argument 'dim' must be int, not {type(value).__name__}"
                 )
 
+        def _invalid_dimname(value):
+            raise RuntimeError(
+                "Invalid name: a valid identifier contains only digits, alphabetical characters, "
+                f"and/or underscore and starts with a non-digit. got: '{value}'."
+            )
+
+        def _validate_all_any_dim(value):
+            if value is None:
+                return
+            if isinstance(value, bool):
+                _raise_invalid_combo_with_got("(Tensor, dim=bool)", {"dim_detail": "bool"})
+                return
+            if isinstance(value, int):
+                return
+            if isinstance(value, str):
+                if value.isidentifier():
+                    raise RuntimeError(
+                        f"{op_short_name}: You passed a dimname (string) to this op in place of a dimension "
+                        "index but it does not yet support this behavior. Please pass a dimension index to "
+                        "work around this."
+                    )
+                _invalid_dimname(value)
+                return
+            if isinstance(value, (list, tuple)):
+                if not value:
+                    return
+                first = value[0]
+                if isinstance(first, (bool, str)):
+                    _raise_invalid_combo_with_got("(Tensor, dim=list)", {"dim_detail": "list"})
+                    return
+                for item in value:
+                    if not isinstance(item, int):
+                        _raise_invalid_combo_with_got("(Tensor, dim=list)", {"dim_detail": "list"})
+                        return
+                return
+            _raise_invalid_combo()
+
+        def _validate_count_nonzero_dim(value):
+            if value is None:
+                return
+            if isinstance(value, bool):
+                _raise_invalid_combo_with_got("(Tensor, dim=bool)", {"dim_detail": "bool"})
+                return
+            if isinstance(value, str):
+                _raise_invalid_combo_with_got("(Tensor, dim=str)", {"dim_detail": "str"})
+                return
+            if isinstance(value, int):
+                return
+            if isinstance(value, (list, tuple)):
+                if not value:
+                    return
+                for item in value:
+                    if not isinstance(item, int) or isinstance(item, bool):
+                        _raise_invalid_combo_with_got("(Tensor, dim=list)", {"dim_detail": "list"})
+                        return
+                return
+            _raise_invalid_combo()
+
         def _type_label(value):
             if isinstance(value, bool):
                 return "bool"
@@ -559,6 +617,12 @@ def _validate_transpose_dims(dim0, dim1):
             if op_short_name in {"argmax", "argmin"} and param.name == "dim":
                 _validate_arg_reduce_dim(value)
                 continue
+            if op_short_name in {"all", "any"} and param.name == "dim":
+                _validate_all_any_dim(value)
+                continue
+            if op_short_name == "count_nonzero" and param.name == "dim":
+                _validate_count_nonzero_dim(value)
+                continue
             if op_short_name == "view" and param.name == "shape":
                 _validate_view_shape(value)
                 continue

src/mindtorch_v2/_dispatch/schemas.py

@@ -236,6 +236,24 @@ def register_schemas():
     registry.register_schema("setitem", "setitem(Tensor(a!) self, Any key, Any value) -> Tensor")
 
     _register_reduction_ops(("all", "any", "argmax", "argmin", "count_nonzero"))
+    registry.register_error_overrides(
+        "all",
+        {
+            "unexpected": "{name}() received an invalid combination of arguments - got {got}, but expected one of:\n * (Tensor input, *, Tensor out = None)\n * (Tensor input, tuple of ints dim = None, bool keepdim = False, *, Tensor out = None)\n * (Tensor input, int dim, bool keepdim = False, *, Tensor out = None)\n * (Tensor input, name dim, bool keepdim = False, *, Tensor out = None)\n",
+        },
+    )
+    registry.register_error_overrides(
+        "any",
+        {
+            "unexpected": "{name}() received an invalid combination of arguments - got {got}, but expected one of:\n * (Tensor input, *, Tensor out = None)\n * (Tensor input, tuple of ints dim = None, bool keepdim = False, *, Tensor out = None)\n * (Tensor input, int dim, bool keepdim = False, *, Tensor out = None)\n * (Tensor input, name dim, bool keepdim = False, *, Tensor out = None)\n",
+        },
+    )
+    registry.register_error_overrides(
+        "count_nonzero",
+        {
+            "unexpected": "{name}() received an invalid combination of arguments - got {got}, but expected one of:\n * (Tensor input, int dim = None)\n      didn't match because some of the keywords were incorrect: dim\n * (Tensor input, tuple of ints dim)\n      didn't match because some of the arguments have invalid types: (Tensor, !dim={dim_detail}!)\n",
+        },
+    )
     registry.register_schema("amin", "amin(Tensor input, int[]? dim=None, bool keepdim=False) -> Tensor")
     registry.register_schema("amax", "amax(Tensor input, int[]? dim=None, bool keepdim=False) -> Tensor")
 

tests/mindtorch_v2/contract/test_schema_dim_validation.py

@@ -519,3 +519,75 @@ def th_call():
         pt.argmin(pt.tensor([1.0, 2.0]), dim="0")
 
     assert_torch_error(mt_call, th_call)
+
+
+def test_dispatch_all_rejects_bool_dim_matches_torch():
+    mt_x = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
+
+    def mt_call():
+        dispatch("all", mt_x.device.type, mt_x, dim=True)
+
+    def th_call():
+        pt.all(pt.tensor([[1.0, 0.0], [0.0, 1.0]]), dim=True)
+
+    assert_torch_error(mt_call, th_call)
+
+
+def test_dispatch_all_rejects_str_dim_matches_torch():
+    mt_x = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
+
+    def mt_call():
+        dispatch("all", mt_x.device.type, mt_x, dim="0")
+
+    def th_call():
+        pt.all(pt.tensor([[1.0, 0.0], [0.0, 1.0]]), dim="0")
+
+    assert_torch_error(mt_call, th_call)
+
+
+def test_dispatch_any_rejects_bool_dim_matches_torch():
+    mt_x = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
+
+    def mt_call():
+        dispatch("any", mt_x.device.type, mt_x, dim=True)
+
+    def th_call():
+        pt.any(pt.tensor([[1.0, 0.0], [0.0, 1.0]]), dim=True)
+
+    assert_torch_error(mt_call, th_call)
+
+
+def test_dispatch_any_rejects_str_dim_matches_torch():
+    mt_x = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
+
+    def mt_call():
+        dispatch("any", mt_x.device.type, mt_x, dim="0")
+
+    def th_call():
+        pt.any(pt.tensor([[1.0, 0.0], [0.0, 1.0]]), dim="0")
+
+    assert_torch_error(mt_call, th_call)
+
+
+def test_dispatch_count_nonzero_rejects_bool_dim_matches_torch():
+    mt_x = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
+
+    def mt_call():
+        dispatch("count_nonzero", mt_x.device.type, mt_x, dim=True)
+
+    def th_call():
+        pt.count_nonzero(pt.tensor([[1.0, 0.0], [0.0, 1.0]]), dim=True)
+
+    assert_torch_error(mt_call, th_call)
+
+
+def test_dispatch_count_nonzero_rejects_str_dim_matches_torch():
+    mt_x = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
+
+    def mt_call():
+        dispatch("count_nonzero", mt_x.device.type, mt_x, dim="0")
+
+    def th_call():
+        pt.count_nonzero(pt.tensor([[1.0, 0.0], [0.0, 1.0]]), dim="0")
+
+    assert_torch_error(mt_call, th_call)