snn.layers

Base

class snn.layers.base.Layer(trainable=True, name=None)

Bases: object

build(input_shape)

forward(x, training=False)

backward(grad)

property params

property grads

get_config()

count_params()

set_param(key, val)

Set a single parameter by name.

Override in composite layers (e.g. TransformerBlock) to route prefixed keys to the correct sub-layer. The default implementation calls setattr(self, key, val) when the attribute exists.

Dense

class snn.layers.Dense(units, activation=None, use_bias=True, kernel_initializer='glorot_uniform', bias_initializer='zeros', name=None)

Bases: Layer

Fully-connected (dense) layer.

Parameters:

• units (int) – Number of output neurons.

• activation (str or Activation, optional) – Activation function applied after the linear transform.

• use_bias (bool) – Whether to include a bias vector.

• kernel_initializer (str or callable) – Initializer for the weight matrix W.

• bias_initializer (str or callable) – Initializer for the bias vector b.

build(input_shape)

forward(x, training=False)

backward(grad)

property params

property grads

get_config()

Conv2D

class snn.layers.Conv2D(filters, kernel_size=3, stride=1, padding='valid', activation=None, use_bias=True, kernel_initializer='he_normal', bias_initializer='zeros', name=None)

Bases: Layer

2D Convolutional layer.

Expects input shape: (N, H, W, C_in) — channels-last convention.

Parameters:

• filters (int) – Number of output filters (channels).

• kernel_size (int or (int, int)) – Size of the convolution kernel.

• stride (int) – Stride of the convolution.

• padding (str or int) – ‘same’, ‘valid’, or an explicit integer pad.

• activation (str or Activation) – Applied after convolution + bias.

• use_bias (bool)

• kernel_initializer (str or callable)

• bias_initializer (str or callable)

build(input_shape)

forward(x, training=False)

backward(grad)

property params

property grads

get_config()

Pooling

class snn.layers.MaxPooling2D(pool_size=2, stride=None, name=None)

Bases: Layer

2D Max Pooling layer (channels-last: N, H, W, C).

forward(x, training=False)

backward(grad)

class snn.layers.AveragePooling2D(pool_size=2, stride=None, name=None)

Bases: Layer

2D Average Pooling layer (channels-last: N, H, W, C).

forward(x, training=False)

backward(grad)

class snn.layers.GlobalAveragePooling2D(trainable=True, name=None)

Bases: Layer

Global Average Pooling — reduces (N, H, W, C) to (N, C).

forward(x, training=False)

backward(grad)

class snn.layers.GlobalMaxPooling2D(trainable=True, name=None)

Bases: Layer

Global Max Pooling — reduces (N, H, W, C) to (N, C).

forward(x, training=False)

backward(grad)

Normalisation

class snn.layers.BatchNormalization(axis=-1, momentum=0.99, epsilon=1e-05, center=True, scale=True, name=None)

Bases: Layer

Batch Normalization layer.

Normalizes activations over the batch dimension. During training uses batch statistics; during inference uses running (exponential moving average) statistics.

Parameters:

• axis (int) – Axis to normalize (usually the features/channels axis). For Dense output (N, F) use axis=-1. For Conv2D output (N, H, W, C) use axis=-1.

• momentum (float) – Momentum for the moving average (typically 0.99).

• epsilon (float) – Small constant for numerical stability.

• center (bool) – If True, add learnable beta offset.

• scale (bool) – If True, multiply by learnable gamma scale.

build(input_shape)

forward(x, training=False)

backward(grad)

property params

property grads

get_config()

class snn.layers.LayerNormalization(epsilon=1e-06, center=True, scale=True, name=None)

Bases: Layer

Layer Normalization — normalizes over the last axis (feature axis).

build(input_shape)

forward(x, training=False)

backward(grad)

property params

property grads

Regularisation

class snn.layers.Dropout(rate=0.5, seed=None, name=None)

Bases: Layer

Dropout regularization layer.

During training, randomly sets a fraction rate of inputs to zero and scales the remaining values by 1 / (1 - rate) (inverted dropout). During inference, the layer is a pass-through.

Parameters:

• rate (float) – Fraction of units to drop (0 <= rate < 1).

• seed (int, optional) – Random seed for reproducibility.

forward(x, training=False)

backward(grad)

get_config()

class snn.layers.SpatialDropout2D(rate=0.5, seed=None, name=None)

Bases: Layer

Spatial Dropout for 2D feature maps (N, H, W, C).

Drops entire feature maps (channels) instead of individual elements.

Parameters:

• rate (float) – Fraction of feature maps to drop.

• seed (int, optional)

forward(x, training=False)

backward(grad)

get_config()

Shape Manipulation

class snn.layers.Flatten(trainable=True, name=None)

Bases: Layer

Flattens input to (N, -1). Keeps the batch dimension.

forward(x, training=False)

backward(grad)

class snn.layers.Reshape(target_shape, name=None)

Bases: Layer

Reshapes input (excluding batch dim) to a target shape.

Parameters:

target_shape (tuple) – New shape for each sample (not including batch size).

forward(x, training=False)

backward(grad)

get_config()

Conv1D

class snn.layers.Conv1D(filters, kernel_size, stride=1, padding='valid', activation=None, use_bias=True, kernel_initializer='glorot_uniform', name=None)

Bases: Layer

1-D Convolution layer (channels-last: batch, length, channels).

Slides a bank of filters kernels of width kernel_size along the time/sequence axis. Useful for local pattern detection in sequences (text, audio, time-series).

Parameters:

• filters (int) – Number of output channels (kernels).

• kernel_size (int) – Width of each convolutional kernel.

• stride (int) – Step between successive windows (default 1).

• padding (str) – "valid" — no padding, output shorter than input. "same" — zero-pad so output length equals input length (only when stride=1).

• activation (str or Activation) – Applied after the convolution.

• use_bias (bool) – Whether to include a bias term.

• kernel_initializer (str) – Weight initialiser (default "glorot_uniform").

• shape (Input / output)

• --------------------

• ``(batch (→) –

• length

• in_channels)``

• ``(batch –

• out_length

• filters)``

• 1. (where out_length = (length + 2*pad - kernel_size) // stride +)

Examples

>>> conv = Conv1D(filters=32, kernel_size=3, activation='relu')
>>> x = np.random.randn(4, 50, 8)   # (batch, seq, channels)
>>> out = conv.forward(x)           # (4, 48, 32) with valid padding

build(input_shape)

forward(x, training=False)

backward(grad)

property params

property grads

get_config()

1D Pooling

class snn.layers.MaxPooling1D(pool_size=2, stride=None, name=None)

Bases: Layer

1-D Max Pooling over the time/sequence axis.

Parameters:

• pool_size (int) – Size of the pooling window (default 2).

• stride (int) – Step between windows (default equals pool_size).

• shape (Input / output)

• --------------------

• ``(batch (channels)`` →) –

• length

• ``(batch –

• out_length

• channels)``

forward(x, training=False)

backward(grad)

class snn.layers.AveragePooling1D(pool_size=2, stride=None, name=None)

Bases: Layer

1-D Average Pooling over the time/sequence axis.

Parameters:

• pool_size (int) – Width of each pooling window (default 2).

• stride (int) – Step between windows (default equals pool_size).

forward(x, training=False)

backward(grad)

class snn.layers.GlobalAveragePooling1D(trainable=True, name=None)

Bases: Layer

Global Average Pooling over the time/sequence axis.

Reduces (batch, seq_len, features) to (batch, features) by averaging over seq_len. Standard way to collapse a Transformer or LSTM output into a fixed-size representation for classification.

Examples

>>> gap = GlobalAveragePooling1D()
>>> x = np.random.randn(8, 20, 64)   # (batch, seq, features)
>>> out = gap.forward(x)             # (8, 64)

forward(x, training=False)

backward(grad)

class snn.layers.GlobalMaxPooling1D(trainable=True, name=None)

Bases: Layer

Global Max Pooling over the time/sequence axis.

Reduces (batch, seq_len, features) to (batch, features) by taking the maximum over seq_len.

Examples

>>> gmp = GlobalMaxPooling1D()
>>> x = np.random.randn(8, 20, 64)
>>> out = gmp.forward(x)    # (8, 64)

forward(x, training=False)

backward(grad)

Recurrent

class snn.layers.SimpleRNN(units, activation='tanh', return_sequences=False, return_state=False, kernel_initializer='glorot_uniform', recurrent_initializer='orthogonal', bias_initializer='zeros', name=None)

Bases: Layer

Vanilla (Elman) Recurrent Neural Network cell.

Parameters:

• units (int) – Number of hidden units.

• activation (str or Activation) – Activation applied to the hidden state (default: ‘tanh’).

• return_sequences (bool) – If True, return output at every time step; else only the last.

• return_state (bool) – If True, return (output, last_hidden_state) tuple.

• kernel_initializer (str or callable)

• recurrent_initializer (str or callable)

• bias_initializer (str or callable)

build(input_shape)

forward(x, training=False)

backward(grad)

property params

property grads

class snn.layers.LSTM(units, return_sequences=False, return_state=False, kernel_initializer='glorot_uniform', recurrent_initializer='orthogonal', bias_initializer='zeros', name=None)

Bases: Layer

Long Short-Term Memory (LSTM) layer.

Implements the standard LSTM equations:

f_t = sigmoid(x_t @ Wf + h_{t-1} @ Uf + bf) i_t = sigmoid(x_t @ Wi + h_{t-1} @ Ui + bi) g_t = tanh(x_t @ Wg + h_{t-1} @ Ug + bg) o_t = sigmoid(x_t @ Wo + h_{t-1} @ Uo + bo) c_t = f_t * c_{t-1} + i_t * g_t h_t = o_t * tanh(c_t)

Parameters:

• units (int)

• return_sequences (bool)

• return_state (bool)

build(input_shape)

forward(x, training=False)

backward(grad)

property params

property grads

class snn.layers.GRU(units, return_sequences=False, return_state=False, kernel_initializer='glorot_uniform', bias_initializer='zeros', name=None)

Bases: Layer

Gated Recurrent Unit (GRU) layer.

Equations:

z_t = sigmoid(x_t @ Wz + h_{t-1} @ Uz + bz) # update gate r_t = sigmoid(x_t @ Wr + h_{t-1} @ Ur + br) # reset gate n_t = tanh(x_t @ Wn + (r_t * h_{t-1}) @ Un + bn) # new gate h_t = (1 - z_t) * n_t + z_t * h_{t-1}

Parameters:

• units (int)

• return_sequences (bool)

• return_state (bool)

build(input_shape)

forward(x, training=False)

backward(grad)

property params

property grads

Language / Attention

class snn.layers.Embedding(vocab_size, embed_dim, embeddings_initializer='random_normal', name=None)

Bases: Layer

Token embedding layer.

Maps integer token indices to dense float vectors. This is the standard first layer for any NLP or sequence model.

Parameters:

• vocab_size (int) – Number of unique tokens (size of the vocabulary).

• embed_dim (int) – Dimensionality of the embedding vectors.

• embeddings_initializer (str) – Initialiser for the embedding matrix (default "random_normal").

• shape (Output)

• -----------

• ``(batch –

• ``[0 (seq_len)`` — integer token indices in) –

• vocab_size)``.

• shape

• ------------

• ``(batch –

• seq_len

• vectors. (embed_dim)`` — dense embedding)

Examples

>>> emb = Embedding(vocab_size=1000, embed_dim=64)
>>> x = np.array([[1, 5, 23, 0], [4, 2, 9, 7]])   # (2, 4) integer tokens
>>> out = emb.forward(x)   # (2, 4, 64)

build(input_shape)

forward(x, training=False)

backward(grad)

property params

property grads

count_params()

get_config()

class snn.layers.PositionalEncoding(max_seq_len=512, name=None)

Bases: Layer

Sinusoidal Positional Encoding (Vaswani et al. 2017).

Adds a fixed, non-learned position signal to a sequence embedding so that the model can distinguish token order. The encoding is defined by:

PE[pos, 2i]   = sin(pos / 10000^(2i/d))
PE[pos, 2i+1] = cos(pos / 10000^(2i/d))

This layer has no trainable parameters — it just adds the PE matrix to the incoming embedding and passes the gradient straight through.

Parameters:

• max_seq_len (int) – Maximum sequence length to pre-compute (default 512).

• shape (Input / output)

• --------------------

• ``(batch –

• seq_len

• in-place). (embed_dim)`` → same shape (PE added)

Examples

>>> pe = PositionalEncoding(max_seq_len=128)
>>> x = np.random.randn(4, 20, 64)    # (batch, seq, embed_dim)
>>> out = pe.forward(x)               # (4, 20, 64) — PE added

forward(x, training=False)

backward(grad)

class snn.layers.MultiHeadAttention(embed_dim, num_heads, dropout_rate=0.0, name=None)

Bases: Layer

Multi-Head Self-Attention (Vaswani et al. 2017).

Computes scaled dot-product attention in parallel over num_heads heads, then projects the concatenated result back to embed_dim.

Attention(Q, K, V) = softmax(Q·Kᵀ / √d_k) · V

Each head operates on a head_dim = embed_dim // num_heads subspace.

Parameters:

• embed_dim (int) – Total embedding dimensionality (must be divisible by num_heads).

• num_heads (int) – Number of parallel attention heads.

• dropout_rate (float) – Dropout applied to the attention weight matrix during training.

• shape (Output)

• -----------

• ``(batch –

• seq_len

• embed_dim)``

• shape

• ------------

• ``(batch –

• seq_len

• embed_dim)``

Examples

>>> attn = MultiHeadAttention(embed_dim=64, num_heads=4)
>>> x = np.random.randn(2, 10, 64)
>>> out = attn.forward(x)   # (2, 10, 64)

build(input_shape)

forward(x, training=False, mask=None)

backward(grad)

property params

property grads

class snn.layers.TransformerBlock(embed_dim, num_heads, ff_dim, dropout_rate=0.1, activation='gelu', name=None)

Bases: Layer

Transformer encoder block (Pre-LayerNorm variant).

Architecture per block:

x → LayerNorm → MultiHeadAttention → Dropout → + x (residual)
  → LayerNorm → FFN(ff_dim) → Dropout → + x (residual)

The FFN consists of two Dense layers: Dense(ff_dim, activation) → Dense(embed_dim).

Pre-LN (applied before each sub-layer) is more numerically stable than Post-LN and typically requires no learning-rate warmup.

Parameters:

• embed_dim (int) – Model dimensionality.

• num_heads (int) – Number of attention heads.

• ff_dim (int) – Hidden size of the feed-forward network (usually 2–4 × embed_dim).

• dropout_rate (float) – Dropout applied inside attention and after each sub-layer.

• activation (str) – Activation used in the FFN hidden layer (default "gelu").

• shape (Input / output)

• --------------------

• ``(batch –

• seq_len

• shape. (embed_dim)`` → same)

Examples

Simple language classifier:

model = Sequential([
    Embedding(vocab_size=5000, embed_dim=64),
    PositionalEncoding(),
    TransformerBlock(embed_dim=64, num_heads=4, ff_dim=128),
    GlobalAveragePooling1D(),
    Dense(2, activation="softmax"),
])

forward(x, training=False)

backward(grad)

property params

property grads

set_param(key, val)

Route a prefixed parameter update to the correct sub-layer.

get_config()

Gated Linear Units

class snn.layers.GLU(units, use_bias=True, kernel_initializer='glorot_uniform', name=None)

Bases: Layer

Gated Linear Unit (Dauphin et al. 2017).

Projects the input to 2 × units, splits into two halves (a, b), and returns a ⊙ σ(b):

z = x @ W + bias          # shape (..., 2*units)
a, b = z[..., :units], z[..., units:]
output = a * sigmoid(b)   # shape (..., units)

This halves the effective hidden size while adding a learnable gating mechanism. Widely used in CNNs and early Transformer variants.

Parameters:

• units (int) – Output dimensionality (half of the internal projection width).

• use_bias (bool) – Whether to add a bias term (default True).

• kernel_initializer (str) – Weight initialiser (default "glorot_uniform").

Examples

>>> glu = GLU(units=32)
>>> x = np.random.randn(4, 20, 64)
>>> out = glu.forward(x)   # (4, 20, 32)

build(input_shape)

forward(x, training=False)

backward(grad)

property params

property grads

get_config()

class snn.layers.SwiGLU(units, use_bias=True, kernel_initializer='glorot_uniform', name=None)

Bases: Layer

SwiGLU — Swish-Gated Linear Unit (Shazeer 2020).

Like GLU but replaces the sigmoid gate with Swish (SiLU):

z = x @ W + bias
a, b = z[..., :units], z[..., units:]
output = a ⊙ swish(b)    where swish(b) = b * σ(b)

SwiGLU is the standard FFN activation in LLaMA, PaLM, Mistral, and Gemma. It outperforms ReLU and GELU on most language tasks despite having 50% more parameters than a plain Dense layer (offset by removing one bias in the original papers).

Parameters:

• units (int) – Output dimensionality.

• use_bias (bool) – Whether to include a bias (default True).

Examples

Replace a standard FFN sub-layer:

# Standard
Dense(ff_dim, activation="gelu") → Dense(embed_dim)

# SwiGLU drop-in (ff_dim must be 2/3 of original for equal params)
SwiGLU(units=ff_dim) → Dense(embed_dim)

build(input_shape)

forward(x, training=False)

backward(grad)

property params

property grads

get_config()

Merge & Skip Connections

class snn.layers.Residual(layers, activation='relu', shortcut=None, name=None)

Bases: Layer

Residual (skip-connection) block — for use inside Sequential.

Wraps a sequence of sub-layers and adds a skip connection:

output = activation(sub_layers(x) + shortcut(x))

If the input and output feature sizes differ, a learnable linear projection is built lazily and used as the shortcut.

Parameters:

• layers (list of Layer) – The main (non-shortcut) path.

• activation (str) – Activation applied after the residual sum (default "relu"). Pass None or "linear" to skip the activation.

• shortcut (Layer, optional) – Explicit shortcut layer. Auto-built if input/output dims differ.

Examples

Basic residual block inside Sequential:

model = Sequential([
    Dense(64, activation='relu'),
    Residual([
        Dense(64, activation='relu'),
        Dense(64),              # no activation — applied after sum
    ]),
    Dense(10, activation='softmax'),
])

Two blocks of different sizes (auto-projection):

model = Sequential([
    Dense(128, activation='relu'),
    Residual([Dense(64, activation='relu'), Dense(64)]),
    # shortcut Dense(64) auto-added because 128 ≠ 64
    Dense(10, activation='softmax'),
])

forward(x, training=False)

backward(grad)

property params

property grads

set_param(key, val)

Set a single parameter by name.

Override in composite layers (e.g. TransformerBlock) to route prefixed keys to the correct sub-layer. The default implementation calls setattr(self, key, val) when the attribute exists.

count_params()

class snn.layers.Add(trainable=True, name=None)

Bases: Layer

Element-wise addition of multiple inputs.

Used in residual / skip-connection architectures. All inputs must have the same shape.

In Sequential context, wrap it in a Residual block. In GraphModel context, call it with a list of Tensor nodes:

h1 = Dense(64, 'relu')(x)
h2 = Dense(64, 'relu')(x)
merged = Add()([h1, h2])

Input / output shape

List of arrays, each (batch, …) → same shape.

forward(inputs, training=False)

backward(grad)

class snn.layers.Concatenate(axis=-1, name=None)

Bases: Layer

Concatenate multiple inputs along one axis.

# GraphModel usage
h1 = Dense(32, 'relu')(x)
h2 = Dense(32, 'relu')(x)
merged = Concatenate(axis=-1)([h1, h2])   # shape (..., 64)

Parameters:

axis (int) – Axis along which to concatenate (default -1).

forward(inputs, training=False)

backward(grad)

class snn.layers.TimeDistributed(layer, name=None)

Bases: Layer

Apply a layer independently to each time step of a sequence.

Reshapes (batch, time, features) → (batch*time, features), runs the wrapped layer, then reshapes back. This lets you apply any layer (Dense, Conv1D, …) to every position in a sequence without writing a Python loop.

Parameters:

• layer (Layer) – The layer to apply at each time step.

• shape (Input / output)

• --------------------

• ``(batch (in_features)`` →) –

• time

• ``(batch –

• time

• out_features)``

Examples

>>> td = TimeDistributed(Dense(32, activation='relu'))
>>> x = np.random.randn(8, 20, 64)
>>> out = td.forward(x)    # (8, 20, 32)

Typical use — dense classification head over RNN output:

model = Sequential([
    LSTM(64, return_sequences=True),         # (batch, seq, 64)
    TimeDistributed(Dense(32, 'relu')),      # (batch, seq, 32)
    TimeDistributed(Dense(vocab_size, 'softmax')),
])

forward(x, training=False)

backward(grad)

property params

property grads

set_param(key, val)

Set a single parameter by name.

Override in composite layers (e.g. TransformerBlock) to route prefixed keys to the correct sub-layer. The default implementation calls setattr(self, key, val) when the attribute exists.

count_params()