nnbma.layers package

Submodules

nnbma.layers.additional_module module

class nnbma.layers.additional_module.AdditionalModule(input_features: int | None, output_features: int | None, device: str = 'cpu')

Bases: Module, ABC

Additional module.

Parameters:

• input_features (int) – Number of input features.

• output_features (int) – Number of output features.

• device (str) – Device to use, by default “cpu”.

abstract forward(x: Tensor) → Tensor

Evaluates the associated pytorch function.

Parameters:

x (Tensor) – Input tensor of shape (?, input_features).

Returns:

Output tensor of shape (?, output_features).

Return type:

Tensor

training: bool

class nnbma.layers.additional_module.AdditionalModuleFromExisting(input_features: int | None, output_features: int | None, operation: Module | Callable, device: str = 'cpu')

Bases: AdditionalModule

Additional module build from an existing Torch module or function. Avoid overriding forward method.

Parameters:

• input_features (int) – Number of input features.

• output_features (int) – Number of output features.

• operation (Module | function) – Operation to apply (Torch module or function).

• device (str) – Device to use, by default “cpu”.

forward(x: Tensor) → Tensor

Evaluates the associated pytorch function.

Parameters:

x (Tensor) – Input tensor of shape (?, input_features).

Returns:

Output tensor of shape (?, output_features).

Return type:

Tensor

training: bool

nnbma.layers.polynomial_expansion module

class nnbma.layers.polynomial_expansion.PolynomialExpansion(n_features: int, order: int, device: str = 'cpu')

Bases: Module

Polynomial expansion layer. For instance, for n_features=3 and order=2,

\[\mathrm{poly}((x_1,\, x_2,\,x_3)) = (x_1,\,x_2,\,x_3,\,x_1^2,\,x_1x_2,\,x_1x_3,\,x_2^2,\,x_2x_3,\,x_3^2)\]

Parameters:

• n_features (int) – Number of input features.

• order (int) – maximum degree of the polynomial expansion.

• device (str, optional) – Device to use, by default “cpu”.

static expanded_features(order: int, n_features: int) → int

Returns the number of augmented polynomial features of order lower or equal to order and of n_features variables.

Parameters:

• order (int) – maximum degree of the polynomial expansion.

• n_features (int) – Number of input features.

Returns:

number of augmented polynomial features of order lower or equal to order and of n_features variables.

Return type:

int

forward(x: Tensor) → Tensor

Applies the polynomial expansion.

Parameters:

x (Tensor) – Input tensor of shape (?, n_features).

Returns:

Output tensor of shape (?, expanded_features(order, n_features)).

Return type:

Tensor

training: bool

update_standardization(x: Tensor | ndarray, reset: bool = False) → None

This optional operation standardizes the output of the layer. Even in case of standardized inputs, the output of this layer are in general not standardized because of the polynomial transformations. This fonction can be called multiple times for different batches. To ensure correct moment calculation, entries must pass this functino only once.

Parameters:

• x (torch.Tensor) – Input batch of shape (?, n_features).

• reset (bool, optional) – If True, reset moment calculation.

nnbma.layers.restrictable_linear module

class nnbma.layers.restrictable_linear.RestrictableLinear(in_features: int, out_features: int, bias: bool = True, device: str | None = None, dtype=None, outputs_names: Sequence[str] | None = None)

Bases: Linear

Restrictable linear layer.

Parameters:

• in_features (int) – input dimension

• out_features (int) – output dimension

• bias (bool, optional) – use a bias vector, by default True

• device (Optional[str], optional) – device on which the graph should be defined (cpu or cuda), by default None

• dtype (_type_, optional) – dtype to be used in computations, by default None

• outputs_names (Optional[Sequence[str]], optional) – sequence of output names, by default None

forward(x: Tensor) → Tensor

Evaluates the linear layer, restricted or not depending of past settings.

Parameters:

x (torch.Tensor) – input tensor of shape (?, in_features)

Returns:

output tensor. This tensor has shape (?, out_features) if the full output set is considered, and with less outputs features in case of restricted output.

Return type:

torch.Tensor

in_features: int

out_features: int

restrict_to_output_subset(indices: Sequence[int]) → None

Restricts the output to an output subset.

Parameters:

indices (Sequence[int]) – index subset to predict.

Raises:

PermissionError – the restriction cannot be applied in train mode. To apply the restriction, first turn the model to .eval() mode.

train(mode: bool = True) → RestrictableLinear

Sets the object in train mode if mode is True and in eval mode else.

Parameters:

mode (bool, optional) – Whether the layer is to be set in train mode (True) or eval mode (False), by default True.

Returns:

The updated object.

Return type:

RestrictableLinear

weight: Tensor

Module contents