Add tensor methods

Author: akhilg-nv

tripy/nvtripy/frontend/ops/_registry.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");
@@ -32,13 +32,23 @@ def register_tensor_method(name: str):
     """
 
     # We make a special exception for "shape" since we actually do want that to be a property
-    allowed_methods = ["shape"]
+    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}"
 
     def impl(func: Callable[..., Any]) -> Callable[..., Any]:
-        TENSOR_METHOD_REGISTRY[name] = func
+        # Don't wrap properties (like shape)
+        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)
+            def method_wrapper(self, *args, **kwargs):
+                return func(self, *args, **kwargs)
+
+            TENSOR_METHOD_REGISTRY[name] = method_wrapper
+
         return func
 
     return impl

tripy/nvtripy/frontend/ops/cast.py

@@ -20,12 +20,14 @@
 from nvtripy.common.datatype import bool as tp_bool
 from nvtripy.common.datatype import float32, int8
 from nvtripy.frontend.ops import utils as op_utils
+from nvtripy.frontend.ops._registry import register_tensor_method
 from nvtripy.frontend.ops.dequantize import dequantize
 from nvtripy.frontend.ops.quantize import quantize
 from nvtripy.trace.ops.cast import Cast
 from nvtripy.utils import wrappers
 
 
+@register_tensor_method("cast")
 @export.public_api(document_under="operations/functions")
 @wrappers.interface(
     dtype_constraints={"input": "T1", "dtype": "T2", wrappers.RETURN_VALUE: "T2"},

tripy/nvtripy/frontend/ops/copy.py

@@ -21,10 +21,12 @@
 from nvtripy.common import device as tp_device
 from nvtripy.common.datatype import DATA_TYPES
 from nvtripy.common.exception import raise_error
+from nvtripy.frontend.ops._registry import register_tensor_method
 from nvtripy.utils import wrappers
 
 
 @export.public_api(document_under="operations/functions")
+@register_tensor_method("copy")
 @wrappers.interface(
     dtype_constraints={"input": "T1", wrappers.RETURN_VALUE: "T1"},
     dtype_variables={"T1": list(DATA_TYPES.keys())},

tripy/nvtripy/frontend/ops/flatten.py

@@ -17,10 +17,12 @@
 from nvtripy import export
 from nvtripy.common.exception import raise_error
 from nvtripy.frontend.ops import utils as op_utils
+from nvtripy.frontend.ops._registry import register_tensor_method
 from nvtripy.utils import wrappers
 
 
 @export.public_api(document_under="operations/functions")
+@register_tensor_method("flatten")
 @wrappers.interface(
     dtype_constraints={"input": "T1", wrappers.RETURN_VALUE: "T1"},
     dtype_variables={"T1": ["float32", "float16", "bfloat16", "int4", "int8", "int32", "int64", "bool"]},

tripy/nvtripy/frontend/ops/permute.py

@@ -20,11 +20,13 @@
 from nvtripy import export
 from nvtripy.common.exception import raise_error
 from nvtripy.frontend.ops import utils as op_utils
+from nvtripy.frontend.ops._registry import register_tensor_method
 from nvtripy.trace.ops.permute import Permute
 from nvtripy.utils import wrappers
 
 
 @export.public_api(document_under="operations/functions")
+@register_tensor_method("permute")
 @wrappers.interface(
     dtype_constraints={"input": "T1", wrappers.RETURN_VALUE: "T1"},
     dtype_variables={"T1": ["float32", "float16", "bfloat16", "int4", "int8", "int32", "int64", "bool"]},

tripy/nvtripy/frontend/ops/reshape.py

@@ -20,6 +20,7 @@
 from nvtripy import export
 from nvtripy.common.exception import raise_error
 from nvtripy.frontend.ops import utils as op_utils
+from nvtripy.frontend.ops._registry import register_tensor_method
 from nvtripy.trace.ops.reshape import Reshape
 from nvtripy.types import ShapeLike
 from nvtripy.utils import wrappers
@@ -43,6 +44,7 @@ def infer_dimensions(input: "nvtripy.Tensor", shape: ShapeLike) -> ShapeLike:
 
 
 @export.public_api(document_under="operations/functions")
+@register_tensor_method("reshape")
 @wrappers.interface(
     dtype_constraints={"input": "T1", wrappers.RETURN_VALUE: "T1"},
     dtype_variables={"T1": ["float32", "float16", "bfloat16", "int4", "int8", "int32", "int64", "bool"]},

tripy/nvtripy/frontend/ops/squeeze.py

@@ -16,10 +16,12 @@
 
 from nvtripy import export, utils
 from nvtripy.frontend.ops import utils as op_utils
+from nvtripy.frontend.ops._registry import register_tensor_method
 from nvtripy.utils import wrappers
 
 
 @export.public_api(document_under="operations/functions")
+@register_tensor_method("squeeze")
 @wrappers.interface(
     dtype_constraints={"input": "T1", wrappers.RETURN_VALUE: "T1"},
     dtype_variables={"T1": ["float32", "float16", "bfloat16", "int8", "int32", "int64", "bool"]},

tripy/nvtripy/frontend/ops/transpose.py

@@ -14,10 +14,12 @@
 # limitations under the License.
 from nvtripy import export
 from nvtripy.common.exception import raise_error
+from nvtripy.frontend.ops._registry import register_tensor_method
 from nvtripy.utils import wrappers
 
 
 @export.public_api(document_under="operations/functions")
+@register_tensor_method("transpose")
 @wrappers.interface(
     dtype_constraints={"input": "T1", wrappers.RETURN_VALUE: "T1"},
     dtype_variables={"T1": ["float32", "float16", "bfloat16", "int8", "int32", "int64", "bool"]},

tripy/nvtripy/frontend/ops/unsqueeze.py

@@ -17,10 +17,12 @@
 
 from nvtripy import export
 from nvtripy.frontend.ops import utils as op_utils
+from nvtripy.frontend.ops._registry import register_tensor_method
 from nvtripy.utils import wrappers
 
 
 @export.public_api(document_under="operations/functions")
+@register_tensor_method("unsqueeze")
 @wrappers.interface(
     dtype_constraints={"input": "T1", wrappers.RETURN_VALUE: "T1"},
     dtype_variables={"T1": ["float32", "float16", "bfloat16", "int8", "int32", "int64", "bool"]},

tripy/nvtripy/frontend/tensor.py

@@ -125,6 +125,11 @@ 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