Linear¶

class nvtripy.Linear(in_features: int, out_features: int, bias: bool = True, dtype: dtype = float32, quant_dtype: dtype | None = None, weight_quant_dim: int | None = None)[source]¶

Bases: Module

Applies a linear transformation to the input:

\(Linear(x) = xW^T + b\)

Parameters:

in_features (int) – Size of input features.
out_features (int) – Size of output features.
bias (Tensor | None) – Whether to include the bias term.
dtype (dtype) – The data type to use for the weight and bias parameters.
quant_dtype (dtype | None) – The data type for quantization.
weight_quant_dim (int | None) – The dimension along which to apply the weight quantization scale.

Example

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)

Local Variables¶

>>> linear
Linear(
    weight: Parameter = (shape=(4, 3), dtype=float32),
    bias: Parameter = (shape=(4,), dtype=float32),
)
>>> linear.state_dict()
{
    weight: tensor(
        [[0, 0, 0],
         [1, 1, 1],
         [2, 2, 2],
         [3, 3, 3]], 
        dtype=float32, loc=gpu:0, shape=(4, 3)),
    bias: tensor([0, 1, 2, 3], dtype=float32, loc=gpu:0, shape=(4,)),
}

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

>>> output
tensor(
    [[0, 1, 2, 3],
     [0, 4, 8, 12]], 
    dtype=float32, loc=gpu:0, shape=(2, 4))

__call__(*args: Any, **kwargs: Any) → Any¶

Calls the module with the specified arguments.

Parameters:

*args (Any) – Positional arguments to the module.
**kwargs (Any) – Keyword arguments to the module.

Returns:

The outputs computed by the module.

Return type:

Any

Example

class Module(tp.Module):
    def forward(self, x):
        return tp.relu(x)


module = Module()

input = tp.arange(-3, 3)
out = module(input)  # Note that we do not call `forward` directly.

Local Variables¶

>>> module
Module(
)
>>> module.state_dict()
{}

>>> input
tensor([-3, -2, -1, 0, 1, 2], dtype=float32, loc=gpu:0, shape=(6,))

>>> out
tensor([0, 0, 0, 0, 1, 2], dtype=float32, loc=gpu:0, shape=(6,))

initialize_dummy_parameters() → None¶

Initializes any uninitialized parameters in the module with dummy values. This is useful for debugging and testing purposes.

Example

linear = tp.Linear(2, 2)
print(linear.state_dict())

linear.initialize_dummy_parameters()
print(linear.state_dict())

Output¶

{'weight': <nvtripy.frontend.module.parameter.DefaultParameter object at 0x7ed8151d7310>, 'bias': <nvtripy.frontend.module.parameter.DefaultParameter object at 0x7ed8151dc8b0>}
{'weight': tensor(
    [[1, 1],
     [1, 1]], 
    dtype=float32, loc=gpu:0, shape=(2, 2)), 'bias': tensor([1, 1], dtype=float32, loc=gpu:0, shape=(2,))}

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

class MyModule(tp.Module):
    def __init__(self):
        super().__init__()
        self.param = tp.ones((2,), dtype=tp.float32)


module = MyModule()

print(f"Before: {module.param}")

module.load_state_dict({"param": tp.zeros((2,), dtype=tp.float32)})

print(f"After: {module.param}")

Output¶

Before: tensor([1, 1], dtype=float32, loc=gpu:0, shape=(2,))
After: tensor([0, 0], dtype=float32, loc=gpu:0, shape=(2,))

See also

state_dict()

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

Returns an iterator over immediate children of this module, yielding tuples containing the name of the child module and the child module itself.

Returns:: An iterator over tuples containing the name of the child module and the child module itself.
Return type:: Iterator[Tuple[str, Module]]

Example

class StackedLinear(tp.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = tp.Linear(2, 2)
        self.linear2 = tp.Linear(2, 2)


stacked_linear = StackedLinear()

for name, module in stacked_linear.named_children():
    print(f"{name}: {type(module).__name__}")

Output¶

linear1: Linear
linear2: 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

class MyModule(tp.Module):
    def __init__(self):
        super().__init__()
        self.alpha = tp.Tensor(1)
        self.beta = tp.Tensor(2)


linear = MyModule()

for name, parameter in linear.named_parameters():
    print(f"{name}: {parameter}")

Output¶

alpha: tensor(1, dtype=int32, loc=cpu:0, shape=())
beta: tensor(2, dtype=int32, loc=cpu: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

class MyModule(tp.Module):
    def __init__(self):
        super().__init__()
        self.param = tp.ones((2,), dtype=tp.float32)
        self.linear1 = tp.Linear(2, 2)
        self.linear2 = tp.Linear(2, 2)


module = MyModule()

state_dict = module.state_dict()

Local Variables¶

>>> state_dict
{
    param: tensor([1, 1], dtype=float32, loc=gpu:0, shape=(2,)),
    linear1.weight: <nvtripy.frontend.module.parameter.DefaultParameter object at 0x7ed815143f10>,
    linear1.bias: <nvtripy.frontend.module.parameter.DefaultParameter object at 0x7ed8151e8f40>,
    linear2.weight: <nvtripy.frontend.module.parameter.DefaultParameter object at 0x7ed81544fd30>,
    linear2.bias: <nvtripy.frontend.module.parameter.DefaultParameter object at 0x7ed81544f070>,
}

dtype: dtype¶: The data type used to perform the operation

weight: Tensor¶: The \(W\) matrix of shape \([\text{out_features}, \text{in_features}]\)

bias: Tensor | None¶: The \(b\) matrix of shape \([\text{out_features},]\)

quant_dtype: dtype | None¶: The quantization data type

weight_quant_dim: int | None¶: The dimension along which to apply the weight quantization scale.

weight_scale: Tensor | None¶: The quantization scale for weight

input_scale: Tensor | None¶: The quantization scale for input

forward(x: Tensor) → Tensor[source]¶

Parameters:: x (Tensor) – The input tensor, of shape \([*, \text{in_features}]\).
Returns:: A tensor of shape \([*, \text{out_features}]\).
Return type:: Tensor