Improve detection of instance methods, test cleanup, simplify method registry

akhilg-nv · akhilg-nv · commit a6e0c3506b23 · 2025-08-07T19:01:20.000-07:00
diff --git a/tripy/nvtripy/frontend/ops/__init__.py b/tripy/nvtripy/frontend/ops/__init__.py
@@ -21,3 +21,13 @@
 import nvtripy.frontend.ops.matmul
 import nvtripy.frontend.ops.shape
 import nvtripy.frontend.ops.slice
+
+# Import regular methods that should be available as tensor methods
+import nvtripy.frontend.ops.cast
+import nvtripy.frontend.ops.copy
+import nvtripy.frontend.ops.reshape
+import nvtripy.frontend.ops.transpose
+import nvtripy.frontend.ops.flatten
+import nvtripy.frontend.ops.permute
+import nvtripy.frontend.ops.squeeze
+import nvtripy.frontend.ops.unsqueeze
diff --git a/tripy/nvtripy/frontend/ops/_registry.py b/tripy/nvtripy/frontend/ops/_registry.py
@@ -29,28 +29,19 @@ def register_tensor_method(name: str):
     Decorator to add the method to the tensor method registry with the name specified.
     This does not use the FunctionRegistry decorator because every tensor method would also be
     registered in the public function registry and we would prefer to avoid having overhead
-    from having to dispatch overloads and check types twice.
+    from having to dispatch overloads and check types twice. This needs to be the top level decorator so we can
+    get input type validation from other decorators like `public_api`.
     """
 
     # We make a special exception for "shape" since we actually do want that to be a property
     # We also add additional methods of the tensor class that are not magic methods
     allowed_methods = ["copy", "cast", "shape", "reshape", "transpose", "flatten", "permute", "squeeze", "unsqueeze"]
     assert name in allowed_methods or name.startswith(
         "__"
-    ), f"The tensor method registry should only be used for magic methods, but was used for: {name}"
+    ), f"The tensor method registry should only be used for magic methods and specially allowed methods, but was used for: {name}"
 
     def impl(func: Callable[..., Any]) -> Callable[..., Any]:
-        if name == "shape":
-            TENSOR_METHOD_REGISTRY[name] = func
-        else:
-            # Create a method wrapper that maps 'self' to the first argument (input)
-            # This is the standard pattern for all tensor methods except 'shape' (which is a property)
-            @wraps(func)
-            def method_wrapper(self, *args, **kwargs):
-                return func(self, *args, **kwargs)
-
-            TENSOR_METHOD_REGISTRY[name] = method_wrapper
-
+        TENSOR_METHOD_REGISTRY[name] = func
         return func
 
     return impl
diff --git a/tripy/nvtripy/frontend/ops/copy.py b/tripy/nvtripy/frontend/ops/copy.py
@@ -25,8 +25,8 @@
 from nvtripy.utils import wrappers
 
 
-@export.public_api(document_under="operations/functions")
 @register_tensor_method("copy")
+@export.public_api(document_under="operations/functions")
 @wrappers.interface(
     dtype_constraints={"input": "T1", wrappers.RETURN_VALUE: "T1"},
     dtype_variables={"T1": list(DATA_TYPES.keys())},
diff --git a/tripy/nvtripy/frontend/ops/flatten.py b/tripy/nvtripy/frontend/ops/flatten.py
@@ -21,8 +21,8 @@
 from nvtripy.utils import wrappers
 
 
-@export.public_api(document_under="operations/functions")
 @register_tensor_method("flatten")
+@export.public_api(document_under="operations/functions")
 @wrappers.interface(
     dtype_constraints={"input": "T1", wrappers.RETURN_VALUE: "T1"},
     dtype_variables={"T1": ["float32", "float16", "bfloat16", "int4", "int8", "int32", "int64", "bool"]},
diff --git a/tripy/nvtripy/frontend/ops/permute.py b/tripy/nvtripy/frontend/ops/permute.py
@@ -25,8 +25,8 @@
 from nvtripy.utils import wrappers
 
 
-@export.public_api(document_under="operations/functions")
 @register_tensor_method("permute")
+@export.public_api(document_under="operations/functions")
 @wrappers.interface(
     dtype_constraints={"input": "T1", wrappers.RETURN_VALUE: "T1"},
     dtype_variables={"T1": ["float32", "float16", "bfloat16", "int4", "int8", "int32", "int64", "bool"]},
diff --git a/tripy/nvtripy/frontend/ops/reshape.py b/tripy/nvtripy/frontend/ops/reshape.py
@@ -43,8 +43,8 @@ def infer_dimensions(input: "nvtripy.Tensor", shape: ShapeLike) -> ShapeLike:
     return {"shape": shape}
 
 
-@export.public_api(document_under="operations/functions")
 @register_tensor_method("reshape")
+@export.public_api(document_under="operations/functions")
 @wrappers.interface(
     dtype_constraints={"input": "T1", wrappers.RETURN_VALUE: "T1"},
     dtype_variables={"T1": ["float32", "float16", "bfloat16", "int4", "int8", "int32", "int64", "bool"]},
diff --git a/tripy/nvtripy/frontend/ops/squeeze.py b/tripy/nvtripy/frontend/ops/squeeze.py
@@ -20,8 +20,8 @@
 from nvtripy.utils import wrappers
 
 
-@export.public_api(document_under="operations/functions")
 @register_tensor_method("squeeze")
+@export.public_api(document_under="operations/functions")
 @wrappers.interface(
     dtype_constraints={"input": "T1", wrappers.RETURN_VALUE: "T1"},
     dtype_variables={"T1": ["float32", "float16", "bfloat16", "int8", "int32", "int64", "bool"]},
diff --git a/tripy/nvtripy/frontend/ops/transpose.py b/tripy/nvtripy/frontend/ops/transpose.py
@@ -18,8 +18,8 @@
 from nvtripy.utils import wrappers
 
 
-@export.public_api(document_under="operations/functions")
 @register_tensor_method("transpose")
+@export.public_api(document_under="operations/functions")
 @wrappers.interface(
     dtype_constraints={"input": "T1", wrappers.RETURN_VALUE: "T1"},
     dtype_variables={"T1": ["float32", "float16", "bfloat16", "int8", "int32", "int64", "bool"]},
diff --git a/tripy/nvtripy/frontend/ops/unsqueeze.py b/tripy/nvtripy/frontend/ops/unsqueeze.py
@@ -21,8 +21,8 @@
 from nvtripy.utils import wrappers
 
 
-@export.public_api(document_under="operations/functions")
 @register_tensor_method("unsqueeze")
+@export.public_api(document_under="operations/functions")
 @wrappers.interface(
     dtype_constraints={"input": "T1", wrappers.RETURN_VALUE: "T1"},
     dtype_variables={"T1": ["float32", "float16", "bfloat16", "int8", "int32", "int64", "bool"]},
diff --git a/tripy/nvtripy/frontend/tensor.py b/tripy/nvtripy/frontend/tensor.py
@@ -125,11 +125,6 @@ def from_trace_tensor(cls, trace_tensor, include_code_index=2):
         return instance
 
     def __getattr__(self, name: str):
-        if name in TENSOR_METHOD_REGISTRY:
-            import types
-
-            return types.MethodType(TENSOR_METHOD_REGISTRY[name], self)
-
         import nvtripy as tp
         from nvtripy.common.exception import search_for_missing_attr
 
diff --git a/tripy/nvtripy/utils/ast.py b/tripy/nvtripy/utils/ast.py
@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -102,9 +102,8 @@ def get_ast_node_func_name(node) -> Optional[str]:
 
 # Gets the column offset of the argument at `index` to function called `func_name` in the provided `code` snippet.
 def get_arg_candidate_column_offsets(
-    code: str, index: int, num_positional: int, func_name: str, is_kwarg: bool, arg_names: List[str]
+    code: str, index: int, num_positional: int, func_name: str, is_kwarg: bool
 ) -> Tuple[int, int]:
-
     candidates = []
 
     result = get_parsed_ast(code)
@@ -123,8 +122,9 @@ def get_arg_candidate_column_offsets(
             if is_kwarg:
                 arg_node = node.keywords[index - num_positional]
             else:
-                # For methods, the `self` argument is omitted from ast.Call.args
-                if "self" in arg_names:
+                # Detect method calls by examining AST structure
+                is_method_call = isinstance(node.func, ast.Attribute)
+                if is_method_call:
                     index -= 1
                 # If the final argument is a starred object, then we treat any args
                 # past the end as pointing to the starred object (this would be a variadic call,
diff --git a/tripy/nvtripy/utils/wrappers.py b/tripy/nvtripy/utils/wrappers.py
@@ -17,6 +17,7 @@
 
 import functools
 import inspect
+import types
 from dataclasses import dataclass
 from textwrap import indent
 from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, Union
@@ -40,7 +41,7 @@ class DataTypeConstraints:
 
 
 # Try to include correct column offsets for non-tensor arguments.
-def _add_column_info(arg, arg_index, is_kwarg, num_positional, func_name, arg_names):
+def _add_column_info(arg, arg_index, is_kwarg, num_positional, func_name):
     from nvtripy.frontend.tensor import Tensor
 
     assert isinstance(arg, Tensor), f"This function should only be called for objects that are already Tensor instances"
@@ -94,7 +95,7 @@ def _add_column_info(arg, arg_index, is_kwarg, num_positional, func_name, arg_na
         dispatch_target = dispatch_target.replace("__r", "__")
 
     candidates = utils.ast.get_arg_candidate_column_offsets(
-        source_info.code, arg_index, num_positional, dispatch_target or func_name, is_kwarg, arg_names
+        source_info.code, arg_index, num_positional, dispatch_target or func_name, is_kwarg
     )
 
     # Only set column range if there is exactly one candidate, otherwise we can't reliably determine
@@ -202,7 +203,6 @@ def add_arg(arg):
                 name in kwargs,
                 len(args),
                 func.__name__,
-                [name for name, _ in merged_args],
             )
 
             dtype = None
diff --git a/tripy/tests/integration/test_cast.py b/tripy/tests/integration/test_cast.py
@@ -45,13 +45,17 @@ class TestCast:
         "input_dtype, target_dtype",
         dtype_pairs,
     )
-    def test_cast(self, input_dtype, target_dtype, eager_or_compiled):
+    @pytest.mark.parametrize("use_tensor_method", [False, True])
+    def test_cast(self, input_dtype, target_dtype, use_tensor_method, eager_or_compiled):
         tp_target_dtype = NUMPY_TO_TRIPY[target_dtype]
 
         # TODO(#222): Integer casts with negative numbers fail in many cases
         input_tensor = tp.copy(tp.Tensor(np.ones((2, 3), dtype=input_dtype)), tp.device("gpu"))
 
-        output = eager_or_compiled(tp.cast, input_tensor, tp_target_dtype)
+        if use_tensor_method:
+            output = eager_or_compiled(lambda t: t.cast(tp_target_dtype), input_tensor)
+        else:
+            output = eager_or_compiled(tp.cast, input_tensor, tp_target_dtype)
 
         np_input = cp.from_dlpack(input_tensor).get()
         assert np.array_equal(cp.from_dlpack(output).get(), np_input.astype(target_dtype))
@@ -85,18 +89,3 @@ def test_cast_from_bool(self, target_dtype, eager_or_compiled):
 
         np_compare_to_zero = np_input.astype(target_dtype) == 0
         assert np.array_equal(tp_compare_to_zero, np_compare_to_zero)
-
-    @pytest.mark.parametrize(
-        "input_dtype, target_dtype",
-        dtype_pairs,
-    )
-    def test_cast_tensor_method(self, input_dtype, target_dtype, eager_or_compiled):
-        """Test that tensor.cast() method works and produces same result as free function."""
-        tp_target_dtype = NUMPY_TO_TRIPY[target_dtype]
-
-        input_tensor = tp.copy(tp.Tensor(np.ones((2, 3), dtype=input_dtype)), tp.device("gpu"))
-
-        output = eager_or_compiled(lambda t: t.cast(tp_target_dtype), input_tensor)
-
-        np_input = cp.from_dlpack(input_tensor).get()
-        assert np.array_equal(cp.from_dlpack(output).get(), np_input.astype(target_dtype))
diff --git a/tripy/tests/integration/test_flatten.py b/tripy/tests/integration/test_flatten.py
@@ -31,10 +31,16 @@ class TestFlatten:
         "shape, start_dim, end_dim, expected_shape",
         test_cases,
     )
-    def test_flatten(self, shape, start_dim, end_dim, expected_shape, eager_or_compiled):
+    @pytest.mark.parametrize("use_tensor_method", [False, True])
+    def test_flatten(self, shape, start_dim, end_dim, expected_shape, use_tensor_method, eager_or_compiled):
         cp_a = cp.arange(np.prod(shape)).reshape(shape).astype(np.float32)
         a = tp.Tensor(cp_a)
-        b = eager_or_compiled(tp.flatten, a, start_dim=start_dim, end_dim=end_dim)
+
+        if use_tensor_method:
+            b = eager_or_compiled(lambda t: t.flatten(start_dim=start_dim, end_dim=end_dim), a)
+        else:
+            b = eager_or_compiled(tp.flatten, a, start_dim=start_dim, end_dim=end_dim)
+
         assert b.shape == expected_shape
         assert np.array_equal(cp.from_dlpack(b).get(), cp_a.reshape(expected_shape).get())
 
@@ -57,15 +63,3 @@ def test_flatten_with_unknown_dims(self, eager_or_compiled):
         a = tp.ones((2, 3, 4, 5))
         b = eager_or_compiled(tp.flatten, a, start_dim=1, end_dim=-1)
         assert np.array_equal(cp.from_dlpack(b).get(), np.ones((2, 60), dtype=np.float32))
-
-    @pytest.mark.parametrize(
-        "shape, start_dim, end_dim, expected_shape",
-        test_cases,
-    )
-    def test_flatten_tensor_method(self, shape, start_dim, end_dim, expected_shape, eager_or_compiled):
-        """Test that tensor.flatten() method works and produces same result as free function."""
-        cp_a = cp.arange(np.prod(shape)).reshape(shape).astype(np.float32)
-        a = tp.Tensor(cp_a)
-        b = eager_or_compiled(lambda t: t.flatten(start_dim=start_dim, end_dim=end_dim), a)
-        assert b.shape == expected_shape
-        assert np.array_equal(cp.from_dlpack(b).get(), cp_a.reshape(expected_shape).get())
diff --git a/tripy/tests/integration/test_reshape.py b/tripy/tests/integration/test_reshape.py
@@ -33,12 +33,16 @@ class TestReshape:
         "shape, new_shape",
         test_cases,
     )
-    def test_static_reshape(self, shape, new_shape, eager_or_compiled):
+    @pytest.mark.parametrize("use_tensor_method", [False, True])
+    def test_static_reshape(self, shape, new_shape, use_tensor_method, eager_or_compiled):
         cp_a = cp.arange(np.prod(shape)).reshape(shape).astype(np.float32)
         a = tp.Tensor(cp_a)
-        b = eager_or_compiled(tp.reshape, a, new_shape)
-        if -1 in new_shape:
-            new_shape = tuple(np.prod(shape) // -np.prod(new_shape) if d == -1 else d for d in new_shape)
+
+        if use_tensor_method:
+            b = eager_or_compiled(lambda t: t.reshape(new_shape), a)
+        else:
+            b = eager_or_compiled(tp.reshape, a, new_shape)
+
         assert np.array_equal(cp.from_dlpack(b).get(), cp_a.reshape(new_shape).get())
 
     def test_reshape_shape_tensor(self, eager_or_compiled):
@@ -51,16 +55,3 @@ def test_reshape_shape_with_unknown(self, eager_or_compiled):
         a = tp.ones((2, 3, 4))
         out = eager_or_compiled(tp.reshape, a, (2, a.shape[1], a.shape[2] / 2, -1))
         assert np.array_equal(cp.from_dlpack(out).get(), np.ones((2, 3, 2, 2), dtype=np.float32))
-
-    @pytest.mark.parametrize(
-        "shape, new_shape",
-        test_cases,
-    )
-    def test_reshape_tensor_method(self, shape, new_shape, eager_or_compiled):
-        """Test that tensor.reshape() method works and produces same result as free function."""
-        cp_a = cp.arange(np.prod(shape)).reshape(shape).astype(np.float32)
-        a = tp.Tensor(cp_a)
-        b = eager_or_compiled(lambda t: t.reshape(new_shape), a)
-        if -1 in new_shape:
-            new_shape = tuple(np.prod(shape) // -np.prod(new_shape) if d == -1 else d for d in new_shape)
-        assert np.array_equal(cp.from_dlpack(b).get(), cp_a.reshape(new_shape).get())
diff --git a/tripy/tests/integration/test_squeeze.py b/tripy/tests/integration/test_squeeze.py
@@ -28,17 +28,13 @@ class TestSqueeze:
         "input_shape, dims, expected_shape",
         test_cases,
     )
-    def test_squeeze(self, input_shape, dims, expected_shape):
+    @pytest.mark.parametrize("use_tensor_method", [False, True])
+    def test_squeeze(self, input_shape, dims, expected_shape, use_tensor_method):
         input_tensor = tp.ones(input_shape, dtype=tp.float32)
-        output_tensor = tp.squeeze(input_tensor, dims=dims)
-        assert output_tensor.shape == expected_shape
 
-    @pytest.mark.parametrize(
-        "input_shape, dims, expected_shape",
-        test_cases,
-    )
-    def test_squeeze_tensor_method(self, input_shape, dims, expected_shape):
-        """Test that tensor.squeeze() method works and produces same result as free function."""
-        input_tensor = tp.ones(input_shape, dtype=tp.float32)
-        output_tensor = input_tensor.squeeze(dims)
+        if use_tensor_method:
+            output_tensor = input_tensor.squeeze(dims)
+        else:
+            output_tensor = tp.squeeze(input_tensor, dims=dims)
+
         assert output_tensor.shape == expected_shape
diff --git a/tripy/tests/integration/test_unsqueeze.py b/tripy/tests/integration/test_unsqueeze.py
@@ -24,31 +24,21 @@
 
 class TestUnsqueezeOp:
     @pytest.mark.parametrize("axis", [-1, 0, 2])
-    def test_unsqueeze_dynamic_op(self, axis, eager_or_compiled):
-        def func(a):
-            return tp.unsqueeze(a, dim=axis)
-
+    @pytest.mark.parametrize("use_tensor_method", [False, True])
+    def test_unsqueeze_dynamic_op(self, axis, use_tensor_method, eager_or_compiled):
         inp = np.ones((4, 2, 2, 3), dtype=np.float32)
 
-        out = eager_or_compiled(func, tp.Tensor(inp))
+        if use_tensor_method:
+            out = eager_or_compiled(lambda t: t.unsqueeze(axis), tp.Tensor(inp))
+        else:
+            out = eager_or_compiled(tp.unsqueeze, tp.Tensor(inp), axis)
+
         ref_out = np.expand_dims(inp, axis=axis)
         assert tp.allclose(out, tp.Tensor(ref_out))
-
         assert out.shape == tuple(ref_out.shape)
 
     def test_unsqueeze_compile(self):
         def func(a):
             return tp.unsqueeze(a, 3) == tp.Tensor(3.0)
 
         c = tp.compile(func, args=[tp.InputInfo(((1, 2, 3), 2, 3), dtype=tp.float32)])
-
-    @pytest.mark.parametrize("axis", [-1, 0, 2])
-    def test_unsqueeze_tensor_method(self, axis, eager_or_compiled):
-        """Test that tensor.unsqueeze() method works and produces same result as free function."""
-        inp = np.ones((4, 2, 2, 3), dtype=np.float32)
-
-        out = eager_or_compiled(lambda t: t.unsqueeze(axis), tp.Tensor(inp))
-        ref_out = np.expand_dims(inp, axis=axis)
-        assert tp.allclose(out, tp.Tensor(ref_out))
-
-        assert out.shape == tuple(ref_out.shape)
diff --git a/tripy/tests/utils/test_ast.py b/tripy/tests/utils/test_ast.py
@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -334,6 +334,6 @@ def test_get_arg_candidate_column_offsets(
             if expected[i] is None:
                 continue
             assert (
-                get_arg_candidate_column_offsets(call_str, i, num_positional, func_name, i >= num_positional, arg_names)
+                get_arg_candidate_column_offsets(call_str, i, num_positional, func_name, i >= num_positional)
                 == expected[i]
             )
