Linear#

API#

class warp_nn.modules.layers.Linear(
in_features: int,
out_features: int,
*,
bias: bool = True,
initialize_parameters: bool = True,
requires_grad: bool = True,
)[source]#

Bases: Module

Apply a linear transformation over the final dimension of the input.

\[\text{Linear}(x) = W \, x + b\]

Learnable parameters:

Name

Shape

Description

\(W\)

weight

(out_features, in_features)

Weights

\(b\)

bias

(out_features, 1)

Bias. Only if bias is true

The parameters are initialized from the uniform distribution \(u(-k, k)\) where \(k = \frac{1}{\sqrt{\text{in\_features}}}\).

Parameters:

  • in_features – The number of input features.

  • out_features – The number of output features.

  • bias – Whether to include a bias term.

  • initialize_parameters – Whether to initialize the parameters with their default/initial values. If false, the parameters are left as uninitialized memory, and reading them (e.g. during a forward pass) before loading their values (e.g. from a state dictionary) is undefined behavior.

  • requires_grad – Whether the parameters and the cached output arrays of the module require gradients.

__call__(
input: array,
) → array[source]#

Forward pass of the module.

Parameters:

input – The input array, with shape (batch_size, in_features).

Returns:

The output array, with shape (batch_size, out_features).