Sequential

class nvtripy.Sequential(*modules: Module | Callable[[Tensor], Tensor] | Dict[str, Module | Callable[[Tensor], Tensor]])[source]

Bases: Module

A module to stack multiple callable layers or modules in a sequential order. The Sequential container can accept either a list of modules/callable objects or a dictionary of named modules/callable objects. Layers are added in the order they are passed, and each is called sequentially during the forward pass.

Parameters:

*modules (Module | Callable[[Tensor], Tensor] | Dict[str, Module | Callable[[Tensor], Tensor]]) – The module(s) or callable(s) to include in the sequence. These must take exactly one input and return exactly one output. Can be passed as individual positional arguments or as a single dictionary of named modules.

Example: Sequential with Positional Arguments
1model = tp.Sequential(tp.Linear(1, 3), tp.Linear(3, 2))
2
3input = tp.Tensor([1.0])
4output = model(input)
Local Variables
>>> model
Sequential(
    0: Module = Linear(
        weight: Parameter = (shape=[3, 1], dtype=float32),
        bias: Parameter = (shape=[3], dtype=float32),
    ),
    1: Module = Linear(
        weight: Parameter = (shape=[2, 3], dtype=float32),
        bias: Parameter = (shape=[2], dtype=float32),
    ),
)
>>> model.state_dict()
{
    0.weight: tensor(
        [[0.0000],
         [1.0000],
         [2.0000]], 
        dtype=float32, loc=gpu:0, shape=(3, 1)),
    0.bias: tensor([0.0000, 1.0000, 2.0000], dtype=float32, loc=gpu:0, shape=(3,)),
    1.weight: tensor(
        [[0.0000, 1.0000, 2.0000],
         [3.0000, 4.0000, 5.0000]], 
        dtype=float32, loc=gpu:0, shape=(2, 3)),
    1.bias: tensor([0.0000, 1.0000], dtype=float32, loc=gpu:0, shape=(2,)),
}

>>> input
tensor([1.0000], dtype=float32, loc=gpu:0, shape=(1,))

>>> output
tensor(
    [[10.0000, 29.0000]], 
    dtype=float32, loc=gpu:0, shape=(1, 2))
Example: Sequential with a Dictionary
1model = tp.Sequential({"layer1": tp.Linear(1, 3), "layer2": tp.Linear(3, 2)})
2
3input = tp.Tensor([1.0])
4output = model(input)
Local Variables
>>> model
Sequential(
    layer1: Module = Linear(
        weight: Parameter = (shape=[3, 1], dtype=float32),
        bias: Parameter = (shape=[3], dtype=float32),
    ),
    layer2: Module = Linear(
        weight: Parameter = (shape=[2, 3], dtype=float32),
        bias: Parameter = (shape=[2], dtype=float32),
    ),
)
>>> model.state_dict()
{
    layer1.weight: tensor(
        [[0.0000],
         [1.0000],
         [2.0000]], 
        dtype=float32, loc=gpu:0, shape=(3, 1)),
    layer1.bias: tensor([0.0000, 1.0000, 2.0000], dtype=float32, loc=gpu:0, shape=(3,)),
    layer2.weight: tensor(
        [[0.0000, 1.0000, 2.0000],
         [3.0000, 4.0000, 5.0000]], 
        dtype=float32, loc=gpu:0, shape=(2, 3)),
    layer2.bias: tensor([0.0000, 1.0000], dtype=float32, loc=gpu:0, shape=(2,)),
}

>>> input
tensor([1.0000], dtype=float32, loc=gpu:0, shape=(1,))

>>> output
tensor(
    [[10.0000, 29.0000]], 
    dtype=float32, loc=gpu:0, shape=(1, 2))
Example: Sequential with Callables
1model = tp.Sequential(
2    tp.Conv(in_channels=2, out_channels=2, kernel_dims=(1, 1), stride=(1, 1)),
3    lambda x: tp.avgpool(x, kernel_dims=(2, 2), stride=(1, 1)),
4)
5
6input = tp.ones((1, 2, 2, 2), dtype=tp.float32)
7output = model(input)
Local Variables
>>> model
Sequential(
    0: Module = Conv(
        bias: Parameter = (shape=[2], dtype=float32),
        weight: Parameter = (shape=[2, 2, 1, 1], dtype=float32),
    ),
)
>>> model.state_dict()
{
    0.bias: tensor([0.0000, 1.0000], dtype=float32, loc=gpu:0, shape=(2,)),
    0.weight: tensor(
        [[[[0.0000]],

          [[1.0000]]],


         [[[2.0000]],

          [[3.0000]]]], 
        dtype=float32, loc=gpu:0, shape=(2, 2, 1, 1)),
}

>>> input
tensor(
    [[[[1.0000, 1.0000],
       [1.0000, 1.0000]],

      [[1.0000, 1.0000],
       [1.0000, 1.0000]]]], 
    dtype=float32, loc=gpu:0, shape=(1, 2, 2, 2))

>>> output
tensor(
    [[[[1.0000]],

      [[6.0000]]]], 
    dtype=float32, loc=gpu:0, shape=(1, 2, 1, 1))
__call__(input: Tensor) → Tensor[source]

Defines the forward pass by applying each module in the container sequentially to input

Parameters:

input (Tensor) – The input tensor to pass through the sequence of modules.

Returns:

The output tensor after passing through each module in sequence.

Return type:

Tensor

load_state_dict(state_dict: Dict[str, Tensor], strict: bool = True) → Tuple[Set[str], Set[str]]

Loads parameters from the provided state_dict into the current module. This will recurse over any nested child modules.

Parameters:

  • state_dict (Dict[str, Tensor]) – A dictionary mapping names to parameters.

  • strict (bool) – If True, keys in state_dict must exactly match those in this module. If not, an error will be raised.

Returns:

  • missing_keys: keys that are expected by this module but not provided in state_dict.

  • unexpected_keys: keys that are not expected by this module but provided in state_dict.

Return type:

A tuple of two sets of strings representing

Example
1# Using the `module` and `state_dict` from the `state_dict()` example:
2print(f"Before: {module.param}")
3
4state_dict["param"] = tp.zeros((2,), dtype=tp.float32)
5module.load_state_dict(state_dict)
6
7print(f"After: {module.param}")
Output
Before: tensor([1.0000, 1.0000], dtype=float32, loc=gpu:0, shape=(2,))
After: tensor([0.0000, 0.0000], dtype=float32, loc=gpu:0, shape=(2,))

See also

state_dict()

named_children() → Iterator[Tuple[str, Module]][source]

Returns an iterator over all the first-order modules in this Sequential container. Each child module is represented by its name and the module object itself.

Returns:

An iterator over tuples containing the name and module of each child.

Return type:

Iterator[Tuple[str, Module]]

Example
1model = tp.Sequential(tp.Linear(1, 3), tp.Linear(3, 2))
2
3for name, child in model.named_children():
4    print(f"{name}: {type(child).__name__}")
Local Variables
>>> model
Sequential(
    0: Module = Linear(
        weight: Parameter = (shape=[3, 1], dtype=float32),
        bias: Parameter = (shape=[3], dtype=float32),
    ),
    1: Module = Linear(
        weight: Parameter = (shape=[2, 3], dtype=float32),
        bias: Parameter = (shape=[2], dtype=float32),
    ),
)
>>> model.state_dict()
{
    0.weight: tensor(
        [[0.0000],
         [1.0000],
         [2.0000]], 
        dtype=float32, loc=gpu:0, shape=(3, 1)),
    0.bias: tensor([0.0000, 1.0000, 2.0000], dtype=float32, loc=gpu:0, shape=(3,)),
    1.weight: tensor(
        [[0.0000, 1.0000, 2.0000],
         [3.0000, 4.0000, 5.0000]], 
        dtype=float32, loc=gpu:0, shape=(2, 3)),
    1.bias: tensor([0.0000, 1.0000], dtype=float32, loc=gpu:0, shape=(2,)),
}

>>> child
Linear(
    weight: Parameter = (shape=[2, 3], dtype=float32),
    bias: Parameter = (shape=[2], dtype=float32),
)
>>> child.state_dict()
{
    weight: tensor(
        [[0.0000, 1.0000, 2.0000],
         [3.0000, 4.0000, 5.0000]], 
        dtype=float32, loc=gpu:0, shape=(2, 3)),
    bias: tensor([0.0000, 1.0000], dtype=float32, loc=gpu:0, shape=(2,)),
}
Output
0: Linear
1: Linear
named_parameters() → Iterator[Tuple[str, Tensor]]
Returns:

An iterator over tuples containing the name of a parameter and the parameter itself.

Return type:

Iterator[Tuple[str, Tensor]]

Example
 1class MyModule(tp.Module):
 2    def __init__(self):
 3        super().__init__()
 4        self.alpha = tp.Tensor(1)
 5        self.beta = tp.Tensor(2)
 6
 7
 8linear = MyModule()
 9
10for name, parameter in linear.named_parameters():
11    print(f"{name}: {parameter}")
Output
alpha: tensor(1, dtype=int32, loc=gpu:0, shape=())
beta: tensor(2, dtype=int32, loc=gpu:0, shape=())
state_dict() → Dict[str, Tensor]

Returns a dictionary mapping names to parameters in the module. This will recurse over any nested child modules.

Returns:

A dictionary mapping names to parameters.

Return type:

Dict[str, Tensor]

Example
 1class MyModule(tp.Module):
 2    def __init__(self):
 3        super().__init__()
 4        self.param = tp.ones((2,), dtype=tp.float32)
 5        self.linear1 = tp.Linear(2, 2)
 6        self.linear2 = tp.Linear(2, 2)
 7
 8
 9module = MyModule()
10
11state_dict = module.state_dict()
Local Variables
>>> state_dict
{
    param: tensor([1.0000, 1.0000], dtype=float32, loc=gpu:0, shape=(2,)),
    linear1.weight: tensor(
        [[0.0000, 1.0000],
         [2.0000, 3.0000]], 
        dtype=float32, loc=gpu:0, shape=(2, 2)),
    linear1.bias: tensor([0.0000, 1.0000], dtype=float32, loc=gpu:0, shape=(2,)),
    linear2.weight: tensor(
        [[0.0000, 1.0000],
         [2.0000, 3.0000]], 
        dtype=float32, loc=gpu:0, shape=(2, 2)),
    linear2.bias: tensor([0.0000, 1.0000], dtype=float32, loc=gpu:0, shape=(2,)),
}