Tripy Testing

tests/
├── frontend/
│   ├── ops/
│   │   ├── test_binary.py
│   │   ├── test_softmax.py
│   │   └── ...
│   ├── module/
│   │   ├── test_linear.py
│   │   ├── test_layernorm.py
│   │   └── ...
│   ├── test_tensor.py
│   └── ...
├── trace/
│   ├── ops/
│   │   ├── test_binary.py
│   │   └── ...
│   └── ...
├── backend/
│   └── ...
├── integration/
│   └── ...
├── helper.py          # Test utilities
└── conftest.py        # Shared fixtures

from tests import helper
import nvtripy as tp

# Basic error test
with helper.raises(tp.TripyException):
    result = bad_operation()

# With message matching (regex)
with helper.raises(tp.TripyException, match="Expected.*one of"):
    result = bad_operation()

# With stack info verification
a = tp.Tensor([1.0])
with helper.raises(tp.TripyException, has_stack_info_for=[a]):
    result = bad_operation(a)

with helper.config("enable_input_validation", False):
    # Validation disabled in this block
    result = operation()
# Automatically restored

from tests.helper import NUMPY_TO_TRIPY

# Maps numpy dtypes to tripy dtypes:
# bool → tp.bool, np.int8 → tp.int8, np.int32 → tp.int32,
# np.int64 → tp.int64, np.float16 → tp.float16, np.float32 → tp.float32

import cupy as cp
import numpy as np
import nvtripy as tp


class TestMyOp:
    def test_basic(self):
        input = tp.Tensor([1.0, 2.0, 3.0])
        output = tp.my_op(input)

        expected = np.array([...])  # Compute expected result
        assert np.array_equal(cp.from_dlpack(output).get(), expected)

    def test_with_specific_dim(self):
        input = tp.iota([2, 3], dtype=tp.float32)
        output = tp.my_op(input, dim=1)

        assert output.shape == [2, 3]

import pytest

class TestMyOp:
    @pytest.mark.parametrize("dtype", [tp.float32, tp.float16, tp.bfloat16])
    def test_supported_dtypes(self, dtype):
        input = tp.ones([2, 3], dtype=dtype)
        output = tp.my_op(input)
        assert output.dtype == dtype

    @pytest.mark.parametrize(
        "dtype",
        [tp.int8, tp.int32],
        ids=["int8", "int32"],
    )
    def test_unsupported_dtypes_fail(self, dtype):
        input = tp.ones([2, 3], dtype=dtype)
        with helper.raises(tp.TripyException, match="Expected.*one of"):
            tp.my_op(input)

from tests.helper import NUMPY_TO_TRIPY

class TestTensor:
    @pytest.mark.parametrize("dtype", list(NUMPY_TO_TRIPY.keys()))
    def test_dtype_from_numpy(self, dtype):
        np_array = np.array([1, 2, 3], dtype=dtype)
        tensor = tp.Tensor(np_array)
        assert tensor.dtype == NUMPY_TO_TRIPY[dtype]

class TestBinaryOps:
    def test_mismatched_dtypes_fails(self):
        a = tp.Tensor([1.0, 2.0])
        b = tp.ones((2,), dtype=tp.float16)
        with helper.raises(tp.TripyException):
            c = a + b

class TestLinear:
    def test_basic(self):
        linear = tp.Linear(3, 4)
        linear.weight = tp.iota(linear.weight.shape)
        linear.bias = tp.iota(linear.bias.shape)

        input = tp.iota((2, 3))
        output = linear(input)

        assert cp.from_dlpack(output).get().shape == (2, 4)

    def test_no_bias(self):
        linear = tp.Linear(3, 4, bias=False)
        linear.weight = tp.iota(linear.weight.shape)

        input = tp.iota((2, 3))
        output = linear(input)
        assert output.shape == [2, 4]

    def test_state_dict(self):
        linear = tp.Linear(3, 4)
        sd = linear.state_dict()
        assert "weight" in sd
        assert "bias" in sd

from nvtripy.trace.ops.my_op import MyOp
from nvtripy.trace.ops.base import TraceOp

class TestMyTraceOp:
    def test_creates_correct_op(self):
        input = tp.Tensor([1.0, 2.0])
        output = tp.my_op(input)

        assert isinstance(output.trace_tensor.producer, MyOp)

    def test_infer_rank(self):
        input = tp.Tensor([1.0, 2.0])
        output = tp.my_op(input)

        assert output.trace_tensor.rank == 1

class TestSoftmax:
    def test_matches_torch(self):
        input = tp.iota([2, 3], dtype=tp.float32)
        output = tp.softmax(input, dim=1)

        torch_input = torch.from_dlpack(input)
        torch_output = torch.softmax(torch_input, dim=1)

        assert tp.allclose(output, tp.Tensor(torch_output))

@pytest.mark.parametrize(
    "tensor_a, tensor_b, rtol, atol, expected",
    [
        (tp.Tensor([1.0]), tp.Tensor([1.0]), 1e-5, 1e-8, True),
        (tp.Tensor([1.0]), tp.Tensor([2.0]), 1e-5, 1e-8, False),
    ],
    ids=["equal", "not_equal"],
)
def test_allclose(self, tensor_a, tensor_b, rtol, atol, expected):
    result = tp.allclose(tensor_a, tensor_b, rtol=rtol, atol=atol)
    assert result == expected

# Run all tests
pytest tests/

# Run specific test file
pytest tests/frontend/ops/test_softmax.py

# Run specific test
pytest tests/frontend/ops/test_softmax.py::TestSoftmax::test_basic

# Run with verbose output
pytest -v tests/frontend/ops/test_softmax.py

# Run tests matching a pattern
pytest -k "softmax"