diff --git a/src/mindtorch_v2/_dispatch/schema.py b/src/mindtorch_v2/_dispatch/schema.py index d76b1c9eb..5eb113c25 100644 --- a/src/mindtorch_v2/_dispatch/schema.py +++ b/src/mindtorch_v2/_dispatch/schema.py @@ -380,6 +380,20 @@ def _validate_nan_reduction_dim(value): f"{op_short_name}(): argument 'dim' must be tuple of ints, not {type(value).__name__}" ) + def _validate_logsumexp_dim(value, input_tensor): + if isinstance(value, bool): + _raise_invalid_combo_with_got("(Tensor, dim=bool)") + return + if isinstance(value, int) and not isinstance(value, bool): + return + if isinstance(value, str): + if value.isidentifier(): + raise _dimname_not_found(value, input_tensor) + 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" @@ -659,6 +673,9 @@ def _validate_transpose_dims(dim0, dim1): if op_short_name in {"nansum", "nanmean"} and param.name == "dim": _validate_nan_reduction_dim(value) continue + if op_short_name == "logsumexp" and param.name == "dim": + _validate_logsumexp_dim(value, bound.get("input")) + 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 d13e36f85..fbe64e37c 100644 --- a/src/mindtorch_v2/_dispatch/schemas.py +++ b/src/mindtorch_v2/_dispatch/schemas.py @@ -511,6 +511,12 @@ def register_schemas(): # New GROUP C ops for Tensor API alignment registry.register_schema("logsumexp", "logsumexp(Tensor input, int dim, bool keepdim=False) -> Tensor") + registry.register_error_overrides( + "logsumexp", + { + "unexpected": "{name}() received an invalid combination of arguments - got {got}, but expected one of:\n * (Tensor input, tuple of ints dim, bool keepdim = False, *, Tensor out = None)\n * (Tensor input, tuple of names dim, bool keepdim = False, *, Tensor out = None)\n", + }, + ) registry.register_schema("trace", "trace(Tensor input) -> Tensor") registry.register_schema("det", "det(Tensor input) -> Tensor") registry.register_schema("matrix_power", "matrix_power(Tensor input, int n) -> Tensor") diff --git a/tests/mindtorch_v2/contract/test_schema_dim_validation.py b/tests/mindtorch_v2/contract/test_schema_dim_validation.py index 54e2f4548..14ea1a64c 100644 --- a/tests/mindtorch_v2/contract/test_schema_dim_validation.py +++ b/tests/mindtorch_v2/contract/test_schema_dim_validation.py @@ -687,3 +687,39 @@ def th_call(): pt.nanmean(pt.tensor([1.0, 0.0]), dim="0") assert_torch_error(mt_call, th_call) + + +def test_dispatch_logsumexp_rejects_bool_dim_matches_torch(): + mt_x = torch.tensor([1.0, 0.0]) + + def mt_call(): + dispatch("logsumexp", mt_x.device.type, mt_x, dim=True) + + def th_call(): + pt.logsumexp(pt.tensor([1.0, 0.0]), dim=True) + + assert_torch_error(mt_call, th_call) + + +def test_dispatch_logsumexp_rejects_invalid_name_dim_matches_torch(): + mt_x = torch.tensor([1.0, 0.0]) + + def mt_call(): + dispatch("logsumexp", mt_x.device.type, mt_x, dim="0") + + def th_call(): + pt.logsumexp(pt.tensor([1.0, 0.0]), dim="0") + + assert_torch_error(mt_call, th_call) + + +def test_dispatch_logsumexp_rejects_missing_name_dim_matches_torch(): + mt_x = torch.tensor([1.0, 0.0]) + + def mt_call(): + dispatch("logsumexp", mt_x.device.type, mt_x, dim="x") + + def th_call(): + pt.logsumexp(pt.tensor([1.0, 0.0]), dim="x") + + assert_torch_error(mt_call, th_call)