From 0d193fa5737b8857b5ee613c6d9cc6c6944555e5 Mon Sep 17 00:00:00 2001
From: lvyufeng <lvyufeng@cqu.edu.cn>
Date: Sun, 8 Mar 2026 14:53:37 +0800
Subject: [PATCH] mindtorch_v2: align nansum and nanmean dim schema errors

---
 src/mindtorch_v2/_dispatch/schema.py          | 19 ++++++++
 .../contract/test_schema_dim_validation.py    | 48 +++++++++++++++++++
 2 files changed, 67 insertions(+)

diff --git a/src/mindtorch_v2/_dispatch/schema.py b/src/mindtorch_v2/_dispatch/schema.py
index aa5475d10..d76b1c9eb 100644
--- a/src/mindtorch_v2/_dispatch/schema.py
+++ b/src/mindtorch_v2/_dispatch/schema.py
@@ -364,6 +364,22 @@ def _validate_std_var_dim(value):
                     f"and/or underscore and starts with a non-digit. got: '{value}'."
                 )
 
+        def _validate_nan_reduction_dim(value):
+            if value is None:
+                return
+            if isinstance(value, int) and not isinstance(value, bool):
+                return
+            if isinstance(value, (list, tuple)):
+                for item in value:
+                    if not isinstance(item, int) or isinstance(item, bool):
+                        raise TypeError(
+                            f"{op_short_name}(): argument 'dim' must be tuple of ints, not {type(item).__name__}"
+                        )
+                return
+            raise TypeError(
+                f"{op_short_name}(): argument 'dim' must be tuple of ints, not {type(value).__name__}"
+            )
+
         def _type_label(value):
             if isinstance(value, bool):
                 return "bool"
@@ -640,6 +656,9 @@ def _validate_transpose_dims(dim0, dim1):
             if op_short_name in {"std", "var"} and param.name == "dim":
                 _validate_std_var_dim(value)
                 continue
+            if op_short_name in {"nansum", "nanmean"} and param.name == "dim":
+                _validate_nan_reduction_dim(value)
+                continue
             if op_short_name == "view" and param.name == "shape":
                 _validate_view_shape(value)
                 continue
diff --git a/tests/mindtorch_v2/contract/test_schema_dim_validation.py b/tests/mindtorch_v2/contract/test_schema_dim_validation.py
index 56ebb9d2b..54e2f4548 100644
--- a/tests/mindtorch_v2/contract/test_schema_dim_validation.py
+++ b/tests/mindtorch_v2/contract/test_schema_dim_validation.py
@@ -639,3 +639,51 @@ def th_call():
         pt.var(pt.tensor([1.0, 0.0]), dim="0")
 
     assert_torch_error(mt_call, th_call)
+
+
+def test_dispatch_nansum_rejects_bool_dim_matches_torch():
+    mt_x = torch.tensor([1.0, 0.0])
+
+    def mt_call():
+        dispatch("nansum", mt_x.device.type, mt_x, dim=True)
+
+    def th_call():
+        pt.nansum(pt.tensor([1.0, 0.0]), dim=True)
+
+    assert_torch_error(mt_call, th_call)
+
+
+def test_dispatch_nansum_rejects_str_dim_matches_torch():
+    mt_x = torch.tensor([1.0, 0.0])
+
+    def mt_call():
+        dispatch("nansum", mt_x.device.type, mt_x, dim="0")
+
+    def th_call():
+        pt.nansum(pt.tensor([1.0, 0.0]), dim="0")
+
+    assert_torch_error(mt_call, th_call)
+
+
+def test_dispatch_nanmean_rejects_bool_dim_matches_torch():
+    mt_x = torch.tensor([1.0, 0.0])
+
+    def mt_call():
+        dispatch("nanmean", mt_x.device.type, mt_x, dim=True)
+
+    def th_call():
+        pt.nanmean(pt.tensor([1.0, 0.0]), dim=True)
+
+    assert_torch_error(mt_call, th_call)
+
+
+def test_dispatch_nanmean_rejects_str_dim_matches_torch():
+    mt_x = torch.tensor([1.0, 0.0])
+
+    def mt_call():
+        dispatch("nanmean", mt_x.device.type, mt_x, dim="0")
+
+    def th_call():
+        pt.nanmean(pt.tensor([1.0, 0.0]), dim="0")
+
+    assert_torch_error(mt_call, th_call)