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
1linear = tp.Linear(3, 4)
2
3input = tp.iota((2, 3))
4output = 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.0000, 1.0000, 2.0000],
         [3.0000, 4.0000, 5.0000],
         [6.0000, 7.0000, 8.0000],
         [9.0000, 10.0000, 11.0000]], 
        dtype=float32, loc=gpu:0, shape=(4, 3)),
    bias: tensor([0.0000, 1.0000, 2.0000, 3.0000], dtype=float32, loc=gpu:0, shape=(4,)),
}

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

>>> output
tensor(
    [[0.0000, 1.0000, 2.0000, 3.0000],
     [3.0000, 13.0000, 23.0000, 33.0000]], 
    dtype=float32, loc=gpu:0, shape=(2, 4))
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]]

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
 1class StackedLinear(tp.Module):
 2    def __init__(self):
 3        super().__init__()
 4        self.linear1 = tp.Linear(2, 2)
 5        self.linear2 = tp.Linear(2, 2)
 6
 7
 8stacked_linear = StackedLinear()
 9
10for name, module in stacked_linear.named_children():
11    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
 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,)),
}
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 \([1, \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

__call__(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