From f6ca883b41e78104e960839eed7568d98556098d Mon Sep 17 00:00:00 2001 From: lvyufeng Date: Sun, 8 Mar 2026 14:26:05 +0800 Subject: [PATCH] mindtorch_v2: align std and var dim schema errors --- src/mindtorch_v2/_dispatch/schema.py | 17 +++++++ src/mindtorch_v2/_dispatch/schemas.py | 12 +++++ .../contract/test_schema_dim_validation.py | 48 +++++++++++++++++++ 3 files changed, 77 insertions(+) diff --git a/src/mindtorch_v2/_dispatch/schema.py b/src/mindtorch_v2/_dispatch/schema.py index 4e5a9d088..aa5475d10 100644 --- a/src/mindtorch_v2/_dispatch/schema.py +++ b/src/mindtorch_v2/_dispatch/schema.py @@ -350,6 +350,20 @@ def _validate_count_nonzero_dim(value): return _raise_invalid_combo() + def _validate_std_var_dim(value): + if value is None: + return + if isinstance(value, bool): + _raise_invalid_combo_with_got("(Tensor, dim=bool)") + return + if isinstance(value, str): + if value.isidentifier(): + raise _dimname_not_found(value, bound.get("input")) + 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 _type_label(value): if isinstance(value, bool): return "bool" @@ -623,6 +637,9 @@ def _validate_transpose_dims(dim0, dim1): if op_short_name == "count_nonzero" and param.name == "dim": _validate_count_nonzero_dim(value) continue + if op_short_name in {"std", "var"} and param.name == "dim": + _validate_std_var_dim(value) + continue if op_short_name == "view" and param.name == "shape": _validate_view_shape(value) continue diff --git a/src/mindtorch_v2/_dispatch/schemas.py b/src/mindtorch_v2/_dispatch/schemas.py index d14245aca..d13e36f85 100644 --- a/src/mindtorch_v2/_dispatch/schemas.py +++ b/src/mindtorch_v2/_dispatch/schemas.py @@ -179,6 +179,12 @@ def register_schemas(): }, ) registry.register_schema("std", "std(Tensor input, int[]? dim=None, bool keepdim=False, bool unbiased=True) -> Tensor") + registry.register_error_overrides( + "std", + { + "unexpected": "{name}() received an invalid combination of arguments - got {got}, but expected one of:\n * (Tensor input, tuple of ints dim, bool unbiased = True, bool keepdim = False, *, Tensor out = None)\n * (Tensor input, tuple of ints dim = None, *, Number correction = None, bool keepdim = False, Tensor out = None)\n * (Tensor input, bool unbiased = True)\n didn't match because some of the keywords were incorrect: dim\n * (Tensor input, tuple of names dim, bool unbiased = True, bool keepdim = False, *, Tensor out = None)\n * (Tensor input, tuple of names dim, *, Number correction = None, bool keepdim = False, Tensor out = None)\n", + }, + ) registry.register_schema("reshape", "reshape(Tensor(a) input, int[] shape) -> Tensor(a)") registry.register_error_overrides( @@ -392,6 +398,12 @@ def register_schemas(): registry.register_schema("scatter_add_", "scatter_add_(Tensor(a!) self, int dim, Tensor index, Tensor src) -> Tensor") registry.register_schema("var", "var(Tensor input, int[]? dim=None, bool unbiased=True, bool keepdim=False) -> Tensor") + registry.register_error_overrides( + "var", + { + "unexpected": "{name}() received an invalid combination of arguments - got {got}, but expected one of:\n * (Tensor input, tuple of ints dim, bool unbiased = True, bool keepdim = False, *, Tensor out = None)\n * (Tensor input, tuple of ints dim = None, *, Number correction = None, bool keepdim = False, Tensor out = None)\n * (Tensor input, bool unbiased = True)\n didn't match because some of the keywords were incorrect: dim\n * (Tensor input, tuple of names dim, bool unbiased = True, bool keepdim = False, *, Tensor out = None)\n * (Tensor input, tuple of names dim, *, Number correction = None, bool keepdim = False, Tensor out = None)\n", + }, + ) registry.register_schema("norm", "norm(Tensor input, Any p=2, int[]? dim=None, bool keepdim=False) -> Tensor") registry.register_error_overrides( "norm", diff --git a/tests/mindtorch_v2/contract/test_schema_dim_validation.py b/tests/mindtorch_v2/contract/test_schema_dim_validation.py index 5b049fb6a..56ebb9d2b 100644 --- a/tests/mindtorch_v2/contract/test_schema_dim_validation.py +++ b/tests/mindtorch_v2/contract/test_schema_dim_validation.py @@ -591,3 +591,51 @@ 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) + + +def test_dispatch_std_rejects_bool_dim_matches_torch(): + mt_x = torch.tensor([1.0, 0.0]) + + def mt_call(): + dispatch("std", mt_x.device.type, mt_x, dim=True) + + def th_call(): + pt.std(pt.tensor([1.0, 0.0]), dim=True) + + assert_torch_error(mt_call, th_call) + + +def test_dispatch_std_rejects_str_dim_matches_torch(): + mt_x = torch.tensor([1.0, 0.0]) + + def mt_call(): + dispatch("std", mt_x.device.type, mt_x, dim="0") + + def th_call(): + pt.std(pt.tensor([1.0, 0.0]), dim="0") + + assert_torch_error(mt_call, th_call) + + +def test_dispatch_var_rejects_bool_dim_matches_torch(): + mt_x = torch.tensor([1.0, 0.0]) + + def mt_call(): + dispatch("var", mt_x.device.type, mt_x, dim=True) + + def th_call(): + pt.var(pt.tensor([1.0, 0.0]), dim=True) + + assert_torch_error(mt_call, th_call) + + +def test_dispatch_var_rejects_str_dim_matches_torch(): + mt_x = torch.tensor([1.0, 0.0]) + + def mt_call(): + dispatch("var", mt_x.device.type, mt_x, dim="0") + + def th_call(): + pt.var(pt.tensor([1.0, 0.0]), dim="0") + + assert_torch_error(mt_call, th_call)