tripy-new-module

Adding a New Module to nvtripy

When to Use

• Creating a new neural network layer (e.g., normalization, attention, convolution)
• Implementing a Module subclass with learnable parameters
• Adding a module that wraps existing ops into a reusable component

Architecture Overview

Modules live in nvtripy/frontend/module/ and follow this pattern:

1. Optional helper function: A standalone function (not exported) that implements the math, decorated with @wrappers.interface for constraints.
2. Module class: A @dataclass subclass of Module with @export.public_api and @constant_fields.
3. Parameters: Use DefaultParameter (must be set before use) or OptionalParameter (can be None).

Procedure

Step 1: Create the Module File

Create nvtripy/frontend/module/<module_name>.py:

from dataclasses import dataclass
from typing import Optional, Sequence, Union

from nvtripy import export, utils
from nvtripy.common import datatype
from nvtripy.frontend import wrappers
from nvtripy.frontend.module.module import Module
from nvtripy.frontend.module.parameter import DefaultParameter, OptionalParameter
from nvtripy.frontend.tensor import Tensor
from nvtripy.frontend.wrappers import constant_fields
from nvtripy.frontend.ops import utils as op_utils

# If needed, import the trace op:
from nvtripy.trace.ops.my_op import MyOp

from nvtripy.frontend.constraints import GetInput, GetReturn, OneOf


# Optional: standalone function with constraints (used by the module's forward())
@wrappers.interface(
    input_requirements=OneOf(GetInput("input").dtype, [datatype.float32, datatype.float16, datatype.bfloat16])
    & (GetInput("weight").dtype == GetInput("input").dtype),
    output_guarantees=GetReturn(0).dtype == GetInput("input").dtype,
)
def my_layer_func(
    input: "nvtripy.Tensor",
    weight: "nvtripy.Tensor",
    bias: "nvtripy.Tensor",
    eps: float,
) -> "nvtripy.Tensor":
    # Implementation using existing ops or create_op
    return op_utils.create_op(MyOp, [input, weight, bias], eps=eps)


@export.public_api(document_under="operations/modules")
@dataclass
@constant_fields(["dtype"])
class MyLayer(Module):
    r"""
    Brief math description of the layer.

    :math:`\text{MyLayer}(x) = f(x, W, b)`
    """

    dtype: datatype.dtype
    r"""The data type used to perform the operation."""

    weight: Tensor
    r"""The weight parameter of shape :math:`[\text{features}]`."""

    bias: Optional[Tensor]
    r"""The bias parameter of shape :math:`[\text{features}]`."""

    eps: float
    """A small value for numerical stability."""

    def __init__(
        self,
        features: int,
        bias: bool = True,
        dtype: datatype.dtype = datatype.float32,
        eps: float = 1e-5,
    ) -> None:
        r"""
        Args:
            features: Size of the feature dimension.
            bias: Whether to include a bias term.
            dtype: The data type for parameters.
            eps: Small constant for numerical stability.

        .. code-block:: python
            :linenos:

            layer = tp.MyLayer(3)

            layer.weight = tp.iota(layer.weight.shape)
            layer.bias = tp.iota(layer.bias.shape)

            input = tp.iota((2, 3), dim=1)
            output = layer(input)

            assert cp.from_dlpack(output).get().shape == (2, 3)
        """
        super().__init__()

        self.dtype = dtype
        self.weight = DefaultParameter((features,), dtype=dtype)

        self.bias = None
        if bias:
            self.bias = DefaultParameter((features,), dtype=dtype)

        self.eps = eps

    def forward(self, x: "nvtripy.Tensor") -> "nvtripy.Tensor":
        r"""
        Args:
            x: The input tensor.

        Returns:
            The output tensor.
        """
        return my_layer_func(x, self.weight, self.bias, self.eps)

Step 2: Understand Parameter Types

DefaultParameter(shape, dtype): Creates a placeholder that MUST be replaced with real data before the module runs. Used for required weights:

self.weight = DefaultParameter((out_features, in_features), dtype=dtype)

OptionalParameter(shape, dtype): Can be None — used for optional weights like quantization scales:

self.input_scale = OptionalParameter(shape=[], dtype=dtype)

Step 3: Use @constant_fields

The @constant_fields(["field1", "field2"]) decorator marks fields as compile-time constants. These fields will be baked into the compiled graph and cannot change at runtime. Use for:

• dtype — data type configuration
• normalized_shape — shape parameters that define the layer structure
• quant_dtype — quantization configuration

Step 4: Register the Module

The @export.public_api(document_under="operations/modules") decorator handles registration. The module will be accessible as tp.MyLayer(...).

Ensure the module file is imported in nvtripy/frontend/module/__init__.py.

Complete Example: LayerNorm

# Helper function with constraints
@wrappers.interface(
    input_requirements=OneOf(GetInput("input").dtype, [datatype.float32, datatype.float16, datatype.bfloat16])
    & (GetInput("weight").dtype == GetInput("input").dtype)
    & (GetInput("bias").dtype == GetInput("input").dtype),
    output_guarantees=GetReturn(0).dtype == GetInput("input").dtype,
)
def layernorm(input, weight, bias, eps):
    normalized_shape = weight.shape
    D = len(normalized_shape)
    input_rank = input.rank

    if input_rank < 2:
        raise_error(f"Input must have rank >= 2, got {input.rank}")

    if input_rank > D:
        broadcast_shape = (1,) * (input_rank - D) + normalized_shape
        weight = reshape(weight, broadcast_shape)
        bias = reshape(bias, broadcast_shape)

    return op_utils.create_op(LayerNormOp, [input, weight, bias],
                               normalized_shape=normalized_shape, eps=eps)


@export.public_api(document_under="operations/modules")
@dataclass
@constant_fields(["dtype", "normalized_shape"])
class LayerNorm(Module):
    dtype: datatype.dtype
    normalized_shape: Sequence[int]
    weight: Tensor
    bias: Tensor
    eps: float

    def __init__(self, normalized_shape, dtype=datatype.float32, eps=1e-5):
        super().__init__()
        self.dtype = dtype
        if isinstance(normalized_shape, int):
            normalized_shape = (normalized_shape,)
        self.normalized_shape = normalized_shape
        self.weight = DefaultParameter(normalized_shape, dtype=dtype)
        self.bias = DefaultParameter(normalized_shape, dtype=dtype)
        self.eps = eps

    def forward(self, x):
        return layernorm(x, self.weight, self.bias, self.eps)

Complete Example: Linear

Key patterns from Linear:

• Weight shape: (out_features, in_features) — transposed in forward()
• Optional bias with sentinel: self.bias = DefaultParameter(...) if bias else None
• Quantization support with OptionalParameter for scales
• Uses @constant_fields(["dtype", "quant_dtype"]) for compile-time config

Module Base Class Features

The Module base class (nvtripy/frontend/module/module.py) provides:

• state_dict(): Recursively collects all Tensor parameters (supports nested modules, lists, dicts)
• load_state_dict(state_dict, strict=True): Loads parameters with shape/dtype validation
• __setattr__: Validates parameter compatibility on assignment
• __call__: Calls forward() — modules are callable like functions

Checklist

• Module file created in nvtripy/frontend/module/
• Inherits from Module, decorated with @dataclass and @export.public_api
• @constant_fields applied for compile-time configuration fields
• __init__ calls super().__init__() and uses DefaultParameter/OptionalParameter
• forward() method implemented
• Docstrings with math notation and working code examples
• Helper function with @wrappers.interface constraints (if applicable)
• Registered in nvtripy/frontend/module/__init__.py
• Tests added in tests/frontend/module/