nerva_jax.layers

Neural network layers used by the MultilayerPerceptron class.

Layers in this module operate on matrices with row layout (each row is a sample). They expose a minimal interface with feedforward, backpropagate and optimize.

Functions

parse_linear_layer(text, D, K, optimizer, ...)

Parse a textual layer spec and create a configured Layer instance.

Classes

ActivationLayer(D, K, act)	Linear layer followed by a pointwise activation function.
BatchNormalizationLayer(D)	Batch normalization layer with per-feature gamma and beta.
Layer()	Base class for layers of a neural network with data in row layout (each row is a sample: shape (N, D)).
LinearLayer(D, K)	Linear layer: Y = X W^T + b.
LogSoftmaxLayer(D, K)	Linear layer followed by log_softmax over the last dimension.
SReLULayer(D, K, act)	Activation layer with SReLU and trainable activation parameters.
SoftmaxLayer(D, K)	Linear layer followed by softmax over the last dimension.

class nerva_jax.layers.Layer

Bases: object

Base class for layers of a neural network with data in row layout (each row is a sample: shape (N, D)).

feedforward(X: jax.numpy.ndarray) → jax.numpy.ndarray

backpropagate(Y: jax.numpy.ndarray, DY: jax.numpy.ndarray) → None

optimize(eta)

class nerva_jax.layers.LinearLayer(D: int, K: int)

Bases: Layer

Linear layer: Y = X W^T + b.

Shapes: X (N, D) -> Y (N, K), W (K, D), b (K,).

feedforward(X: jax.numpy.ndarray) → jax.numpy.ndarray

backpropagate(Y: jax.numpy.ndarray, DY: jax.numpy.ndarray) → None

input_size() → int

output_size() → int

set_optimizer(optimizer: str)

set_weights(weight_initializer)

class nerva_jax.layers.ActivationLayer(D: int, K: int, act: ActivationFunction)

Bases: LinearLayer

Linear layer followed by a pointwise activation function.

feedforward(X: jax.numpy.ndarray) → jax.numpy.ndarray

backpropagate(Y: jax.numpy.ndarray, DY: jax.numpy.ndarray) → None

class nerva_jax.layers.SReLULayer(D: int, K: int, act: SReLUActivation)

Bases: ActivationLayer

Activation layer with SReLU and trainable activation parameters.

In addition to W and b, this layer optimizes SReLU’s (al, tl, ar, tr).

backpropagate(Y: jax.numpy.ndarray, DY: jax.numpy.ndarray) → None

set_optimizer(optimizer: str)

class nerva_jax.layers.SoftmaxLayer(D: int, K: int)

Bases: LinearLayer

Linear layer followed by softmax over the last dimension.

feedforward(X: jax.numpy.ndarray) → jax.numpy.ndarray

backpropagate(Y: jax.numpy.ndarray, DY: jax.numpy.ndarray) → None

class nerva_jax.layers.LogSoftmaxLayer(D: int, K: int)

Bases: LinearLayer

Linear layer followed by log_softmax over the last dimension.

feedforward(X: jax.numpy.ndarray) → jax.numpy.ndarray

backpropagate(Y: jax.numpy.ndarray, DY: jax.numpy.ndarray) → None

class nerva_jax.layers.BatchNormalizationLayer(D: int)

Bases: Layer

Batch normalization layer with per-feature gamma and beta.

Normalizes inputs across the batch using per-feature statistics. Shapes: X (N, D) -> Y (N, D), gamma/beta (D,).

feedforward(X: jax.numpy.ndarray) → jax.numpy.ndarray

backpropagate(Y: jax.numpy.ndarray, DY: jax.numpy.ndarray) → None

Compute gradients for gamma/beta and propagate DX through BN.

input_size() → int

output_size() → int

set_optimizer(optimizer: str)

nerva_jax.layers.parse_linear_layer(text: str, D: int, K: int, optimizer: str, weight_initializer: str) → Layer

Parse a textual layer spec and create a configured Layer instance.

Supports Linear, Softmax, LogSoftmax, activation names (e.g. ReLU), and SReLU(…). The optimizer and weight initializer are applied.