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