Add tests

Author: akhilg-nv

tripy/tests/integration/test_cast.py

@@ -23,25 +23,27 @@
 from tests.conftest import skip_if_older_than_sm89
 from tests.helper import NUMPY_TO_TRIPY
 
+dtype_pairs = [
+    (np.int32, np.float32),
+    (np.float32, np.int32),
+    (np.int32, np.int8),
+    (np.float32, np.int8),
+    (np.int8, np.int32),
+    (np.int8, np.float32),
+    # important to test conversion into bool because default StableHLO semantics
+    # are simply to truncate to i1, which is not desirable
+    (np.float32, bool),
+    (np.int32, bool),
+    # requires a dequantization first
+    # TODO(#219): Dequantize fails with dynamic shapes
+    # (np.int8, bool),
+]
+
 
 class TestCast:
     @pytest.mark.parametrize(
         "input_dtype, target_dtype",
-        [
-            (np.int32, np.float32),
-            (np.float32, np.int32),
-            (np.int32, np.int8),
-            (np.float32, np.int8),
-            (np.int8, np.int32),
-            (np.int8, np.float32),
-            # important to test conversion into bool because default StableHLO semantics
-            # are simply to truncate to i1, which is not desirable
-            (np.float32, bool),
-            (np.int32, bool),
-            # requires a dequantization first
-            # TODO(#219): Dequantize fails with dynamic shapes
-            # (np.int8, bool),
-        ],
+        dtype_pairs,
     )
     def test_cast(self, input_dtype, target_dtype, eager_or_compiled):
         tp_target_dtype = NUMPY_TO_TRIPY[target_dtype]
@@ -83,3 +85,18 @@ 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))

tripy/tests/integration/test_copy.py

@@ -16,23 +16,40 @@
 import numpy as np
 import nvtripy as tp
 
+import pytest
+
 
 class TestCopy:
-    def test_to_cpu(self):
+    @pytest.mark.parametrize(
+        "copy_func",
+        [
+            lambda tensor, device: tp.copy(tensor, device),  # Free function
+            lambda tensor, device: tensor.copy(device),  # Tensor method
+        ],
+    )
+    def test_copy_tensor_method(self, copy_func):
+        """Test that both copy methods work with compilation."""
         gpu_tensor = tp.Tensor(cp.ones((2, 2), dtype=cp.float32))
         assert gpu_tensor.device.kind == "gpu"
 
-        cpu_tensor = tp.copy(gpu_tensor, tp.device("cpu"))
-        assert cpu_tensor.device.kind == "cpu"
+        cpu_tensor = copy_func(gpu_tensor, tp.device("cpu"))
 
+        assert cpu_tensor.device.kind == "cpu"
         # If the tensor is really in CPU memory, we should be able to construct a NumPy array from it
         assert np.from_dlpack(cpu_tensor).shape == (2, 2)
 
-    def test_to_gpu(self):
+    @pytest.mark.parametrize(
+        "copy_func",
+        [
+            lambda tensor, device: tp.copy(tensor, device),  # Free function
+            lambda tensor, device: tensor.copy(device),  # Tensor method
+        ],
+    )
+    def test_to_gpu(self, copy_func):
         cpu_tensor = tp.Tensor(np.ones((2, 2), dtype=np.float32))
         assert cpu_tensor.device.kind == "cpu"
 
-        gpu_tensor = tp.copy(cpu_tensor, tp.device("gpu"))
+        gpu_tensor = copy_func(cpu_tensor, tp.device("gpu"))
         assert gpu_tensor.device.kind == "gpu"
 
         # If the tensor is really in GPU memory, we should be able to construct a Cupy array from it

tripy/tests/integration/test_flatten.py

@@ -17,17 +17,19 @@
 import pytest
 import nvtripy as tp
 
+test_cases = [
+    ((2, 3, 4), 0, -1, (24,)),  # Flatten all dimensions
+    ((2, 3, 4), 1, -1, (2, 12)),  # Flatten dimensions 1 through end
+    ((2, 3, 4), 1, 2, (2, 12)),  # Flatten dimensions 1 through 2
+    ((2, 3, 4), 0, 1, (6, 4)),  # Flatten dimensions 0 through 1
+    ((2, 3, 4, 5), 1, 3, (2, 60)),  # Flatten dimensions 1 through 3
+]
+
 
 class TestFlatten:
     @pytest.mark.parametrize(
         "shape, start_dim, end_dim, expected_shape",
-        [
-            ((2, 3, 4), 0, -1, (24,)),  # Flatten all dimensions
-            ((2, 3, 4), 1, -1, (2, 12)),  # Flatten dimensions 1 through end
-            ((2, 3, 4), 1, 2, (2, 12)),  # Flatten dimensions 1 through 2
-            ((2, 3, 4), 0, 1, (6, 4)),  # Flatten dimensions 0 through 1
-            ((2, 3, 4, 5), 1, 3, (2, 60)),  # Flatten dimensions 1 through 3
-        ],
+        test_cases,
     )
     def test_flatten(self, shape, start_dim, end_dim, expected_shape, eager_or_compiled):
         cp_a = cp.arange(np.prod(shape)).reshape(shape).astype(np.float32)
@@ -55,3 +57,15 @@ 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())

tripy/tests/integration/test_reshape.py

@@ -20,16 +20,18 @@
 import pytest
 import nvtripy as tp
 
+test_cases = [
+    ((2, 4), (1, 8)),
+    ((2, 4, 8, 9), (8, 8, 9)),
+    ((2, 4), (8,)),  # change rank of output
+    ((2, 4), (1, -1)),  # check negative dim
+]
+
 
 class TestReshape:
     @pytest.mark.parametrize(
         "shape, new_shape",
-        [
-            ((2, 4), (1, 8)),
-            ((2, 4, 8, 9), (8, 8, 9)),
-            ((2, 4), (8,)),  # change rank of output
-            ((2, 4), (1, -1)),  # check negative dim
-        ],
+        test_cases,
     )
     def test_static_reshape(self, shape, new_shape, eager_or_compiled):
         cp_a = cp.arange(np.prod(shape)).reshape(shape).astype(np.float32)
@@ -49,3 +51,16 @@ 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())

tripy/tests/integration/test_squeeze.py

@@ -1,4 +1,4 @@
-# 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");
@@ -15,18 +15,30 @@
 import nvtripy as tp
 import pytest
 
+test_cases = [
+    ((1, 2, 1), 0, (2, 1)),  # Squeeze first dimension
+    ((1, 2, 1), (0, 2), (2,)),  # Squeeze first and third dimensions
+    ((1, 2, 1), tuple(), (1, 2, 1)),  # No dimensions to squeeze
+    ((1, 2, 1), (-3, -1), (2,)),  # Squeeze using negative dimensions
+]
+
 
 class TestSqueeze:
     @pytest.mark.parametrize(
         "input_shape, dims, expected_shape",
-        [
-            ((1, 2, 1), 0, (2, 1)),  # Squeeze first dimension
-            ((1, 2, 1), (0, 2), (2,)),  # Squeeze first and third dimensions
-            ((1, 2, 1), tuple(), (1, 2, 1)),  # No dimensions to squeeze
-            ((1, 2, 1), (-3, -1), (2,)),  # Squeeze using negative dimensions
-        ],
+        test_cases,
     )
     def test_squeeze(self, input_shape, dims, expected_shape):
         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)
+        assert output_tensor.shape == expected_shape

tripy/tests/integration/test_unsqueeze.py

@@ -41,3 +41,14 @@ 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)