apache--tvm/tests/python/relax/test_base_py_module_printer.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=missing-docstring, invalid-name, unused-argument
# ruff: noqa: F401, F841

import pytest

import tvm
from tvm.relax.base_py_module import BasePyModule
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T


@I.ir_module
class SimplePyFuncModule(BasePyModule):
    """Test simple Python functions with basic operations."""

    @I.pyfunc
    def add(self, x, y):
        """Simple addition function."""
        x_tvm = self._convert_pytorch_to_tvm(x)
        y_tvm = self._convert_pytorch_to_tvm(y)
        result = self.call_tir(self.add_tir, [x_tvm, y_tvm], out_ty=R.Tensor((5,), "float32"))
        return self._convert_tvm_to_pytorch(result)

    @I.pyfunc
    def multiply(self, x, y):
        """Simple multiplication function."""
        x_tvm = self._convert_pytorch_to_tvm(x)
        y_tvm = self._convert_pytorch_to_tvm(y)
        result = self.call_tir(self.multiply_tir, [x_tvm, y_tvm], out_ty=R.Tensor((5,), "float32"))
        return self._convert_tvm_to_pytorch(result)

    @T.prim_func(s_tir=True)
    def add_tir(var_x: T.handle, var_y: T.handle, var_out: T.handle):
        x = T.match_buffer(var_x, (5,), "float32")
        y = T.match_buffer(var_y, (5,), "float32")
        out = T.match_buffer(var_out, (5,), "float32")

        for i in range(5):
            out[i] = x[i] + y[i]

    @T.prim_func(s_tir=True)
    def multiply_tir(var_x: T.handle, var_y: T.handle, var_out: T.handle):
        x = T.match_buffer(var_x, (5,), "float32")
        y = T.match_buffer(var_y, (5,), "float32")
        out = T.match_buffer(var_out, (5,), "float32")

        for i in range(5):
            out[i] = x[i] * y[i]

    @R.function
    def main_relax(x: R.Tensor((5,), "float32"), y: R.Tensor((5,), "float32")) -> R.Tensor(
        (5,), "float32"
    ):
        return R.add(x, y)


@I.ir_module
class ComplexPyFuncModule(BasePyModule):
    """Test complex Python logic with ML pipeline and error handling."""

    @I.pyfunc
    def ml_pipeline(self, input_data, model_params):
        """Complex ML pipeline with data validation and error handling."""
        # Data validation
        if input_data is None or model_params is None:
            raise ValueError("Inputs cannot be None")

        try:
            # Convert to TVM format
            tvm_data = self._convert_pytorch_to_tvm(input_data)
            tvm_params = self._convert_pytorch_to_tvm(model_params)

            # Run ML inference
            features = self.call_tir(
                self.extract_features, [tvm_data], out_ty=R.Tensor((10,), "float32")
            )

            predictions = self.call_tir(
                self.ml_inference, [features, tvm_params], out_ty=R.Tensor((5,), "float32")
            )

            # Post-process results
            final_result = self.call_tir(
                self.post_process, [predictions], out_ty=R.Tensor((5,), "float32")
            )

            return self._convert_tvm_to_pytorch(final_result)

        except Exception as e:
            self._log_error(f"ML pipeline failed: {e}")
            return self._get_default_value()

    @I.pyfunc
    def data_preprocessing(self, raw_data):
        """Data preprocessing with conditional logic."""
        if hasattr(raw_data, "numpy"):
            # Vectorized path for numpy-compatible data
            data_np = raw_data.numpy()
            processed = self._vectorized_preprocess(data_np)
        else:
            # Fallback path for other data types
            processed = self._elementwise_preprocess(raw_data)

        # Convert and return
        tvm_processed = self._convert_pytorch_to_tvm(processed)
        result = self.call_tir(
            self.normalize_data, [tvm_processed], out_ty=R.Tensor((10,), "float32")
        )
        return self._convert_tvm_to_pytorch(result)

    @T.prim_func(s_tir=True)
    def extract_features(data: T.handle, features: T.handle):
        T.func_attr({"tirx.noalias": True})
        Data = T.match_buffer(data, (10,), "float32")
        Features = T.match_buffer(features, (10,), "float32")

        for i in range(10):
            Features[i] = T.sqrt(Data[i])

    @T.prim_func(s_tir=True)
    def ml_inference(features: T.handle, params: T.handle, output: T.handle):
        T.func_attr({"tirx.noalias": True})
        Features = T.match_buffer(features, (10,), "float32")
        Params = T.match_buffer(params, (10,), "float32")
        Output = T.match_buffer(output, (5,), "float32")

        for i in range(5):
            Output[i] = Features[i] * Params[i] + Features[i + 5] * Params[i + 5]

    @T.prim_func(s_tir=True)
    def post_process(predictions: T.handle, final: T.handle):
        T.func_attr({"tirx.noalias": True})
        Predictions = T.match_buffer(predictions, (5,), "float32")
        Final = T.match_buffer(final, (5,), "float32")

        for i in range(5):
            Final[i] = T.max(Predictions[i], 0.0)

    @T.prim_func(s_tir=True)
    def normalize_data(data: T.handle, normalized: T.handle):
        T.func_attr({"tirx.noalias": True})
        Data = T.match_buffer(data, (10,), "float32")
        Normalized = T.match_buffer(normalized, (10,), "float32")

        for i in range(10):
            Normalized[i] = Data[i] / 255.0


@I.ir_module
class EdgeCasePyFuncModule(BasePyModule):
    """Test edge cases and boundary conditions."""

    @I.pyfunc
    def empty_func(self):
        """Empty function with no operations."""
        pass

    @I.pyfunc
    def single_return(self, x):
        """Function with immediate return."""
        return x

    @I.pyfunc
    def nested_conditionals(self, data, threshold):
        """Function with complex nested conditional logic."""
        if data is None:
            return None

        if hasattr(data, "shape"):
            if len(data.shape) == 1:
                if data.shape[0] > threshold:
                    return self._process_large_data(data)
                else:
                    return self._process_small_data(data)
            elif len(data.shape) == 2:
                return self._process_2d_data(data)
            else:
                return self._process_nd_data(data)
        else:
            return self._process_scalar_data(data)

    @I.pyfunc
    def loop_with_break(self, data, max_iter):
        """Function with loop and break statement."""
        result = []
        for i, item in enumerate(data):
            if i >= max_iter:
                break
            if item > 0:
                result.append(item * 2)
            else:
                result.append(0)
        return result

    @T.prim_func(s_tir=True)
    def dummy_tir(data: T.handle, output: T.handle):
        T.func_attr({"tirx.noalias": True})
        Data = T.match_buffer(data, (1,), "float32")
        Output = T.match_buffer(output, (1,), "float32")
        Output[0] = Data[0]


@I.ir_module
class PerformancePyFuncModule(BasePyModule):
    """Test performance optimization patterns."""

    @I.pyfunc
    def vectorized_operation(self, x, y):
        """Vectorized operation with numpy fallback."""
        try:
            # Try vectorized operation first
            if hasattr(x, "numpy") and hasattr(y, "numpy"):
                x_np = x.numpy()
                y_np = y.numpy()
                result_np = x_np + y_np
                return self._convert_numpy_to_pytorch(result_np)
        except Exception:
            pass

        # Fallback to TVM processing
        x_tvm = self._convert_pytorch_to_tvm(x)
        y_tvm = self._convert_pytorch_to_tvm(y)
        result = self.call_tir(
            self.vectorized_add, [x_tvm, y_tvm], out_ty=R.Tensor((10,), "float32")
        )
        return self._convert_tvm_to_pytorch(result)

    @I.pyfunc
    def batch_processing(self, batch_data):
        """Batch processing with memory optimization."""
        batch_size = len(batch_data)
        results = []

        # Process in chunks to optimize memory usage
        chunk_size = min(batch_size, 100)
        for i in range(0, batch_size, chunk_size):
            chunk = batch_data[i : i + chunk_size]
            chunk_result = self._process_chunk(chunk)
            results.extend(chunk_result)

        return results

    @I.pyfunc
    def memory_efficient_transform(self, large_tensor):
        """Memory-efficient tensor transformation."""
        # Use in-place operations when possible
        if hasattr(large_tensor, "requires_grad") and not large_tensor.requires_grad:
            # In-place operation for efficiency
            large_tensor.add_(1.0)
            return large_tensor
        else:
            # Create new tensor if gradients are needed
            return large_tensor + 1.0

    @T.prim_func(s_tir=True)
    def vectorized_add(a: T.handle, b: T.handle, c: T.handle):
        T.func_attr({"tirx.noalias": True})
        A = T.match_buffer(a, (10,), "float32")
        B = T.match_buffer(b, (10,), "float32")
        C = T.match_buffer(c, (10,), "float32")

        for i in range(10):
            C[i] = A[i] + B[i]


@I.ir_module
class IntegrationPyFuncModule(BasePyModule):
    """Test integration with external libraries and complex workflows."""

    @I.pyfunc
    def sklearn_integration(self, input_data, scaler_params):
        """Integration with scikit-learn preprocessing."""
        try:
            # Import sklearn components
            from sklearn.decomposition import PCA
            from sklearn.preprocessing import StandardScaler

            # Create and fit scaler
            scaler = StandardScaler()
            if scaler_params is not None:
                scaler.mean_ = scaler_params["mean"]
                scaler.scale_ = scaler_params["scale"]
            else:
                scaler.fit(input_data)

            # Transform data
            scaled_data = scaler.transform(input_data)

            # Apply PCA if needed
            if input_data.shape[1] > 10:
                pca = PCA(n_components=10)
                reduced_data = pca.fit_transform(scaled_data)
            else:
                reduced_data = scaled_data

            # Convert to TVM and process
            tvm_data = self._convert_pytorch_to_tvm(reduced_data)
            result = self.call_tir(
                self.final_transform,
                [tvm_data],
                out_ty=R.Tensor((reduced_data.shape[0], 10), "float32"),
            )

            return self._convert_tvm_to_pytorch(result)

        except ImportError:
            # Fallback if sklearn is not available
            return self._fallback_preprocessing(input_data)

    @I.pyfunc
    def multi_stage_pipeline(self, raw_input):
        """Multi-stage processing pipeline."""
        # Stage 1: Data cleaning
        cleaned = self._clean_data(raw_input)

        # Stage 2: Feature extraction
        features = self._extract_features(cleaned)

        # Stage 3: Model inference
        predictions = self._run_inference(features)

        # Stage 4: Post-processing
        final_result = self._post_process_output(predictions)

        return final_result

    @T.prim_func(s_tir=True)
    def final_transform(data: T.handle, output: T.handle):
        T.func_attr({"tirx.noalias": True})
        Data = T.match_buffer(data, (10, 10), "float32")
        Output = T.match_buffer(output, (10, 10), "float32")

        for i in range(10):
            for j in range(10):
                Output[i, j] = T.tanh(Data[i, j])


@I.ir_module
class ErrorHandlingPyFuncModule(BasePyModule):
    """Test comprehensive error handling and validation."""

    @I.pyfunc
    def robust_data_processing(self, input_data, config):
        """Robust data processing with comprehensive error handling."""
        try:
            # Validate inputs
            if not self._validate_inputs(input_data, config):
                raise ValueError("Invalid input data or configuration")

            # Check data types
            if not self._check_data_types(input_data):
                raise TypeError("Unsupported data types")

            # Process data with retry logic
            max_retries = config.get("max_retries", 3)
            for attempt in range(max_retries):
                try:
                    result = self._process_with_validation(input_data, config)
                    if self._validate_output(result):
                        return result
                    else:
                        raise RuntimeError("Output validation failed")
                except Exception as e:
                    if attempt == max_retries - 1:
                        raise
                    self._log_warning(f"Attempt {attempt + 1} failed: {e}")
                    continue

        except Exception as e:
            self._log_error(f"Data processing failed: {e}")
            return self._get_safe_fallback(input_data, config)

    @I.pyfunc
    def graceful_degradation(self, primary_input, fallback_input):
        """Function that gracefully degrades when primary path fails."""
        try:
            # Try primary processing path
            result = self._primary_processing(primary_input)
            return result
        except Exception as e:
            self._log_warning(f"Primary processing failed: {e}")

            try:
                # Try fallback path
                result = self._fallback_processing(fallback_input)
                return result
            except Exception as e2:
                self._log_error(f"Fallback processing also failed: {e2}")
                # Return safe default
                return self._get_safe_default()

    @T.prim_func(s_tir=True)
    def safe_transform(data: T.handle, output: T.handle):
        T.func_attr({"tirx.noalias": True})
        Data = T.match_buffer(data, (5,), "float32")
        Output = T.match_buffer(output, (5,), "float32")

        for i in range(5):
            # Safe operation that handles edge cases
            if Data[i] > 0:
                Output[i] = T.sqrt(Data[i])
            else:
                Output[i] = 0.0


# Pytest test functions to verify the classes work correctly
def test_simple_pyfunc_module_creation():
    """Test that SimplePyFuncModule can be created."""
    # Get the IRModule instance from the TVMScript decorated class
    ir_mod = SimplePyFuncModule
    device = tvm.cpu()

    # Create BasePyModule instance
    module = BasePyModule(ir_mod, device)
    assert isinstance(module, BasePyModule)

    # Note: Python functions are stored in pyfuncs, not as direct attributes
    # We need to check if they exist in the IRModule's pyfuncs
    if hasattr(ir_mod, "pyfuncs"):
        assert "add" in ir_mod.pyfuncs
        assert "multiply" in ir_mod.pyfuncs

    # Check that TIR functions exist
    assert hasattr(module, "add_tir")
    assert hasattr(module, "multiply_tir")

    # Note: This particular TVMScript is for testing purpose only, and cannot compile
    # Relax functions may not be available due to TVMScript compilation issues
    print("Note: This TVMScript is for testing purpose only, and cannot compile")


def test_complex_pyfunc_module_creation():
    """Test that ComplexPyFuncModule can be created."""
    ir_mod = ComplexPyFuncModule
    device = tvm.cpu()

    module = BasePyModule(ir_mod, device)
    assert isinstance(module, BasePyModule)

    # Check Python functions in pyfuncs
    if hasattr(ir_mod, "pyfuncs"):
        assert "ml_pipeline" in ir_mod.pyfuncs
        assert "data_preprocessing" in ir_mod.pyfuncs

    # Check TIR functions
    assert hasattr(module, "extract_features")
    assert hasattr(module, "ml_inference")
    assert hasattr(module, "post_process")
    assert hasattr(module, "normalize_data")


def test_edge_case_pyfunc_module_creation():
    """Test that EdgeCasePyFuncModule can be created."""
    ir_mod = EdgeCasePyFuncModule
    device = tvm.cpu()

    module = BasePyModule(ir_mod, device)
    assert isinstance(module, BasePyModule)

    # Check Python functions in pyfuncs
    if hasattr(ir_mod, "pyfuncs"):
        assert "empty_func" in ir_mod.pyfuncs
        assert "single_return" in ir_mod.pyfuncs
        assert "nested_conditionals" in ir_mod.pyfuncs
        assert "loop_with_break" in ir_mod.pyfuncs

    # Check TIR function
    assert hasattr(module, "dummy_tir")


def test_performance_pyfunc_module_creation():
    """Test that PerformancePyFuncModule can be created."""
    ir_mod = PerformancePyFuncModule
    device = tvm.cpu()

    module = BasePyModule(ir_mod, device)
    assert isinstance(module, BasePyModule)

    # Check Python functions in pyfuncs
    if hasattr(ir_mod, "pyfuncs"):
        assert "vectorized_operation" in ir_mod.pyfuncs
        assert "batch_processing" in ir_mod.pyfuncs
        assert "memory_efficient_transform" in ir_mod.pyfuncs

    # Check TIR function
    assert hasattr(module, "vectorized_add")


def test_integration_pyfunc_module_creation():
    """Test that IntegrationPyFuncModule can be created."""
    ir_mod = IntegrationPyFuncModule
    device = tvm.cpu()

    module = BasePyModule(ir_mod, device)
    assert isinstance(module, BasePyModule)

    # Check Python functions in pyfuncs
    if hasattr(ir_mod, "pyfuncs"):
        assert "sklearn_integration" in ir_mod.pyfuncs
        assert "multi_stage_pipeline" in ir_mod.pyfuncs

    # Check TIR function
    assert hasattr(module, "final_transform")


def test_error_handling_pyfunc_module_creation():
    """Test that ErrorHandlingPyFuncModule can be created."""
    ir_mod = ErrorHandlingPyFuncModule
    device = tvm.cpu()

    module = BasePyModule(ir_mod, device)
    assert isinstance(module, BasePyModule)

    # Check Python functions in pyfuncs
    if hasattr(ir_mod, "pyfuncs"):
        assert "robust_data_processing" in ir_mod.pyfuncs
        assert "graceful_degradation" in ir_mod.pyfuncs

    # Check TIR function
    assert hasattr(module, "safe_transform")


def test_all_modules_inherit_from_base():
    """Test that all modules properly inherit from BasePyModule."""
    modules = [
        SimplePyFuncModule,
        ComplexPyFuncModule,
        EdgeCasePyFuncModule,
        PerformancePyFuncModule,
        IntegrationPyFuncModule,
        ErrorHandlingPyFuncModule,
    ]

    device = tvm.cpu()
    for ir_mod in modules:
        module = BasePyModule(ir_mod, device)
        assert isinstance(module, BasePyModule)
        assert hasattr(module, "script")
        assert hasattr(module, "show")


def test_pyfunc_decorators():
    """Test that all @I.pyfunc decorated functions are present."""
    ir_mod = SimplePyFuncModule
    device = tvm.cpu()
    module = BasePyModule(ir_mod, device)

    # Check that the functions exist in pyfuncs
    if hasattr(ir_mod, "pyfuncs"):
        assert "add" in ir_mod.pyfuncs
        assert "multiply" in ir_mod.pyfuncs

        # Get the actual function objects
        add_func = ir_mod.pyfuncs["add"]
        multiply_func = ir_mod.pyfuncs["multiply"]

        # Check that they are callable
        assert callable(add_func)
        assert callable(multiply_func)

        # Check function signatures
        import inspect

        add_sig = inspect.signature(add_func)
        assert len(add_sig.parameters) == 3  # self, x, y

        multiply_sig = inspect.signature(multiply_func)
        assert len(multiply_sig.parameters) == 3  # self, x, y


def test_tir_functions():
    """Test that TIR functions are properly defined."""
    ir_mod = SimplePyFuncModule
    device = tvm.cpu()
    module = BasePyModule(ir_mod, device)

    # Check TIR function attributes
    assert hasattr(module, "add_tir")
    assert hasattr(module, "multiply_tir")

    # These should be callable (though they're TIR functions)
    assert callable(module.add_tir)
    assert callable(module.multiply_tir)


def test_relax_functions():
    """Test that Relax functions are properly defined."""
    ir_mod = SimplePyFuncModule
    device = tvm.cpu()
    module = BasePyModule(ir_mod, device)

    # Note: This particular TVMScript is for testing purpose only, and cannot compile
    # Relax functions may not be available due to TVMScript compilation issues
    print("Note: This TVMScript is for testing purpose only, and cannot compile")

    # We can still check that the module was created successfully
    assert isinstance(module, BasePyModule)
    assert hasattr(module, "script")
    assert hasattr(module, "show")


def test_module_docstrings():
    """Test that all modules have proper docstrings."""
    modules = [
        SimplePyFuncModule,
        ComplexPyFuncModule,
        EdgeCasePyFuncModule,
        PerformancePyFuncModule,
        IntegrationPyFuncModule,
        ErrorHandlingPyFuncModule,
    ]

    for module_class in modules:
        # TVMScript decorator changes the class, so we check that it's callable
        # and can create instances instead of checking docstrings
        assert callable(module_class)
        # We can't directly instantiate TVMScript decorated classes
        # but we can create BasePyModule instances with them
        device = tvm.cpu()
        instance = BasePyModule(module_class, device)
        assert isinstance(instance, BasePyModule)


def test_python_function_complexity():
    """Test that complex Python functions have the expected structure."""
    ir_mod = ComplexPyFuncModule
    device = tvm.cpu()
    module = BasePyModule(ir_mod, device)

    # Check that complex functions exist in pyfuncs
    if hasattr(ir_mod, "pyfuncs"):
        assert "ml_pipeline" in ir_mod.pyfuncs
        assert "data_preprocessing" in ir_mod.pyfuncs

        # Get the actual function objects
        ml_func = ir_mod.pyfuncs["ml_pipeline"]
        preprocess_func = ir_mod.pyfuncs["data_preprocessing"]

        # These should be callable
        assert callable(ml_func)
        assert callable(preprocess_func)

        # Check function signatures
        import inspect

        ml_sig = inspect.signature(ml_func)
        assert len(ml_sig.parameters) == 3  # self, input_data, model_params

        preprocess_sig = inspect.signature(preprocess_func)
        assert len(preprocess_sig.parameters) == 2  # self, raw_data


def test_script_and_show_methods():
    """Test that script() and show() methods work correctly."""
    ir_mod = SimplePyFuncModule
    device = tvm.cpu()
    module = BasePyModule(ir_mod, device)

    # Test script() method
    script_output = module.script()
    assert isinstance(script_output, str)
    assert len(script_output) > 0

    # Test show() method
    try:
        module.show()
        # If we get here, show() worked
        assert True
    except Exception as e:
        # If show() fails, the feature is not working properly
        pytest.fail(f"show() method failed: {e}")


def test_python_functions_in_irmodule():
    """Test that Python functions are properly stored in IRModule pyfuncs."""
    ir_mod = SimplePyFuncModule
    device = tvm.cpu()
    module = BasePyModule(ir_mod, device)

    # Check that pyfuncs attribute exists and contains our functions
    if hasattr(ir_mod, "pyfuncs"):
        pyfuncs = ir_mod.pyfuncs
        assert isinstance(pyfuncs, dict)
        assert "add" in pyfuncs
        assert "multiply" in pyfuncs

        # Check that the functions are callable
        assert callable(pyfuncs["add"])
        assert callable(pyfuncs["multiply"])

        # Check function names
        assert pyfuncs["add"].__name__ == "add"
        assert pyfuncs["multiply"].__name__ == "multiply"
    else:
        pytest.fail("pyfuncs attribute not found in IRModule")


def test_call_py_func_with_base_py_module():
    """Test R.call_py_func with BasePyModule."""
    import numpy as np
    import torch

    from tvm.relax import TensorType, Var
    from tvm.relax.expr import StringImm
    from tvm.relax.op import call_py_func

    # Test 1: Operator creation and basic properties
    x = Var("x", TensorType((5,), "float32"))
    y = Var("y", TensorType((5,), "float32"))

    call_expr = call_py_func(StringImm("test_func"), (x, y), out_ty=R.Tensor((5,), "float32"))

    assert call_expr.op.name == "relax.call_py_func"
    assert call_expr.args[0].value == "test_func"
    assert len(call_expr.args) == 2

    # Test 2: Compilation validation
    try:
        call_py_func(
            "invalid",
            (Var("x", TensorType((5,), "float32")),),
            out_ty=R.Tensor((5,), "float32"),
        )
        assert False, "Should raise type error"
    except Exception as e:
        assert "Mismatched type" in str(e) or "Expected" in str(e)

    # Test 3: Validation and error handling
    @I.ir_module
    class ValidationTestModule(BasePyModule):
        @R.function
        def test_invalid_call(x: R.Tensor((5,), "float32")) -> R.Tensor((5,), "float32"):
            result = R.call_py_func("non_existent_func", (x,), out_ty=R.Tensor((5,), "float32"))
            return result

    device = tvm.cpu()
    module = ValidationTestModule(device)

    x = torch.randn(5, dtype=torch.float32)

    with pytest.raises(ValueError, match="Python function 'non_existent_func' not found"):
        module.call_py_func("non_existent_func", [x])

    # Test 4: Using call_py_func within Relax functions
    @I.ir_module
    class RelaxCallPyFuncModule(BasePyModule):
        @I.pyfunc
        def torch_relu(self, x):
            """PyTorch ReLU implementation."""
            return torch.relu(x)

        @I.pyfunc
        def torch_softmax(self, x, dim=0):
            """PyTorch softmax implementation."""
            return torch.softmax(x, dim=dim)

        @R.function
        def mixed_computation(x: R.Tensor((10,), "float32")) -> R.Tensor((10,), "float32"):
            relu_result = R.call_py_func("torch_relu", (x,), out_ty=R.Tensor((10,), "float32"))
            final_result = R.call_py_func(
                "torch_softmax", (relu_result,), out_ty=R.Tensor((10,), "float32")
            )
            return final_result

    device = tvm.cpu()
    module = RelaxCallPyFuncModule(device)

    x = torch.randn(10, dtype=torch.float32)

    expected = torch.softmax(torch.relu(x), dim=0)

    relu_result = module.call_py_func("torch_relu", [x])
    final_result = module.call_py_func("torch_softmax", [relu_result])

    # Convert to numpy for comparison
    if isinstance(final_result, tvm.runtime.Tensor):
        final_result_np = final_result.numpy()
    else:
        final_result_np = final_result

    if isinstance(expected, torch.Tensor):
        expected_np = expected.numpy()
    else:
        expected_np = expected

    # Use numpy for comparison since we have numpy arrays
    tvm.testing.assert_allclose(final_result_np, expected_np, rtol=1e-5, atol=1e-5)


if __name__ == "__main__":
    tvm.testing.main()