Linear

class tripy.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 (Parameter | 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
Example
1linear = tp.Linear(3, 4)
2
3input = tp.iota((2, 3))
4output = linear(input)

>>> 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))
dtype: dtype

The data type used to perform the operation

weight: Parameter

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

bias: Parameter | 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: Parameter | None

The quantization scale for weight

input_scale: Parameter | 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