Operators for Online Learning

Operators for Online Learning#

ETP operators are the first layer of an online-learning-ready network. They mark the parameterized computations that the compiler should connect to temporal hidden state. This chapter introduces the five public operators, then verifies their physical-unit and JAX-transformation contracts.

See the complete ETP Operators API for signatures and generated reference pages.

import brainstate
import brainunit as u
import jax
import jax.numpy as jnp

import braintrace

Function reference#

braintrace provides five user-facing ETP operator functions:

Function

Underlying primitives

Purpose

braintrace.matmul()

etp_mm_p (batched) / etp_mv_p (unbatched)

Dense matrix multiplication

braintrace.element_wise()

etp_elemwise_p

Element-wise (diagonal) weight ops

braintrace.conv()

etp_conv_p

Convolution

braintrace.sparse_matmul()

etp_sp_mm_p / etp_sp_mv_p

Sparse matrix multiplication

braintrace.lora_matmul()

etp_lora_mm_p / etp_lora_mv_p

LoRA (Low-Rank Adaptation) matmul

Each function auto-dispatches between batched and unbatched variants based on input dimensionality. The generated pages above are indexed together in the ETP Operators API.

1. braintrace.matmul(x, weight, bias=None) – Dense Matrix Multiplication#

Computes \(y = x \, @ \, w \; (+ b)\).

Auto-dispatches based on x.ndim:

  • x.ndim >= 2 –> etp_mm_p (batched): expects x of shape (batch, in_features)

  • x.ndim == 1 –> etp_mv_p (unbatched): expects x of shape (in_features,)

# Batched matmul: x has shape (batch, in_features)
x_batched = jnp.ones((4, 3))    # batch=4, in_features=3
w = jnp.ones((3, 5))            # in_features=3, out_features=5

y_batched = braintrace.matmul(x_batched, w)
print("Batched output shape:", y_batched.shape)   # (4, 5)

# Unbatched matmul: x has shape (in_features,)
x_single = jnp.ones((3,))       # in_features=3

y_single = braintrace.matmul(x_single, w)
print("Unbatched output shape:", y_single.shape)   # (5,)

An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.

Batched output shape: (4, 5)
Unbatched output shape: (5,)

# With bias
b = jnp.zeros((5,))

y_with_bias = braintrace.matmul(x_batched, w, bias=b)
print("With bias:", y_with_bias.shape)              # (4, 5)

With bias: (4, 5)

2. braintrace.element_wise(weight, *, weight_fn=None) — Element-wise Operation#

Applies weight_fn to the weight and passes the result through a marker primitive. The operation is treated as diagonal in the hidden-state space.

\[y = \texttt{weight\_fn}(w)\]

weight_fn defaults to None (identity); supply any JAX-differentiable function when you want a non-trivial transformation.

Common use cases:

  • Gating mechanisms in RNNs (learnable gate biases)

  • Learnable time constants or thresholds in spiking neural networks

  • Any parameter that enters the computation element-wise

Note: etp_elemwise_p is the only primitive registered with gradient_enabled=True. The compiler descends into it when walking y -> h, so it does not act as a tail boundary for upstream ETP weights. See the gradient_enabled Flag section below for details.

# Identity (default weight_fn): just marks the weight for ETP
w_gate = jnp.array([0.5, -0.3, 0.8, 0.1])

y_identity = braintrace.element_wise(w_gate)
print("Identity:", y_identity)

# With a transformation function
y_sigmoid = braintrace.element_wise(w_gate, weight_fn=jax.nn.sigmoid)
print("Sigmoid:", y_sigmoid)

# With absolute value (e.g., enforcing positive time constants)
y_abs = braintrace.element_wise(w_gate, weight_fn=jnp.abs)
print("Abs:", y_abs)

Identity: [ 0.5 -0.3  0.8  0.1]
Sigmoid: [0.62245935 0.4255575  0.6899745  0.5249792 ]
Abs: [0.5 0.3 0.8 0.1]

3. braintrace.conv(x, kernel, bias=None, *, strides, padding, ...) – Convolution#

ETP-aware convolution that wraps jax.lax.conv_general_dilated. Computes:

\[y = \text{conv}(x, \text{kernel}) \; (+ b)\]

Important: Always expects a batch dimension on x.

Supports all parameters of jax.lax.conv_general_dilated: strides, padding, lhs_dilation, rhs_dilation, feature_group_count, batch_group_count, and dimension_numbers.

# 1D convolution example
# x: (batch, spatial, channels) with dimension_numbers
x_1d = jnp.ones((2, 16, 3))         # batch=2, length=16, in_channels=3
kernel_1d = jnp.ones((4, 3, 8))     # kernel_size=4, in_channels=3, out_channels=8

y_conv = braintrace.conv(
    x_1d, kernel_1d,
    strides=(1,),
    padding='SAME',
    dimension_numbers=('NWC', 'WIO', 'NWC'),
)
print("Conv1D output shape:", y_conv.shape)  # (2, 16, 8)

Conv1D output shape: (2, 16, 8)

# 2D convolution example
x_2d = jnp.ones((2, 32, 32, 3))          # batch=2, H=32, W=32, in_channels=3
kernel_2d = jnp.ones((3, 3, 3, 16))      # kH=3, kW=3, in_channels=3, out_channels=16

y_conv2d = braintrace.conv(
    x_2d, kernel_2d,
    strides=(1, 1),
    padding='SAME',
    dimension_numbers=('NHWC', 'HWIO', 'NHWC'),
)
print("Conv2D output shape:", y_conv2d.shape)  # (2, 32, 32, 16)

Conv2D output shape: (2, 32, 32, 16)

4. braintrace.sparse_matmul(x, weight_data, *, sparse_mat, bias=None) – Sparse Matmul#

ETP-aware sparse matrix multiplication. Computes:

\[y = x \, @ \, \text{sparse}(w) \; (+ b)\]

The sparse_mat argument provides the sparse structure (indices), while weight_data contains only the non-zero values. This is useful for models with sparse connectivity patterns, such as biologically plausible neural networks or graph neural networks.

import brainevent

# Create a reproducible sparse connectivity matrix
brainstate.random.seed(13)
dense_w = jnp.where(
    brainstate.random.uniform(size=(50, 50)) < 0.1,
    brainstate.random.normal(size=(50, 50)),
    0.0
)
# sparse_mat must be a brainevent.DataRepresentation (e.g. brainevent.CSR),
# which implements the with_data / dt2t_transposed / dt2t protocol.
sparse_mat = brainevent.CSR.fromdense(dense_w)

# The learnable parameter is just the non-zero data
weight_data = sparse_mat.data

x_sp = jnp.ones((4, 50))  # batch=4, features=50
y_sp = braintrace.sparse_matmul(x_sp, weight_data, sparse_mat=sparse_mat)
print("Sparse matmul output shape:", y_sp.shape)  # (4, 50)

5. braintrace.lora_matmul(x, B, A, *, alpha=1.0, bias=None) – LoRA Matmul#

Low-Rank Adaptation matmul. Computes:

\[y = \alpha \cdot x \, @ \, B \, @ \, A \; (+ b)\]

where \(B \in \mathbb{R}^{\text{in} \times \text{rank}}\) and \(A \in \mathbb{R}^{\text{rank} \times \text{out}}\) are low-rank factors. This is useful for parameter-efficient fine-tuning of large models, where only the low-rank factors are trained.

in_features, out_features, rank = 64, 32, 4

brainstate.random.seed(17)
B = brainstate.random.normal(size=(in_features, rank)) * 0.01
A = brainstate.random.normal(size=(rank, out_features)) * 0.01

x_lora = jnp.ones((8, in_features))  # batch=8

y_lora = braintrace.lora_matmul(x_lora, B, A, alpha=2.0)
print("LoRA output shape:", y_lora.shape)  # (8, 32)
print("LoRA output (first sample):", y_lora[0])

Physical Units (brainunit / Quantity) Support#

Every user-facing ETP function accepts brainunit.Quantity inputs. The wrapper separates mantissas and units, binds the JAX primitive to plain arrays, and recombines the result with u.maybe_decimal. Bias values must be dimensionally compatible with the combined input and weight unit.

# Quantity-valued inputs pass through unchanged.
x_q = jnp.ones((4, 3)) * u.volt          # shape (4, 3), unit = V
w_q = jnp.ones((3, 5)) * u.siemens       # shape (3, 5), unit = S
b_q = jnp.zeros((5,)) * u.amp             # must match V * S = A

y_q = braintrace.matmul(x_q, w_q, bias=b_q)
print("Output:", y_q)
print("Unit :", u.get_unit(y_q))

Output: [[3. 3. 3. 3. 3.]
 [3. 3. 3. 3. 3.]
 [3. 3. 3. 3. 3.]
 [3. 3. 3. 3. 3.]] A
Unit : A

JAX Compatibility#

ETP operators participate in standard JAX transformations such as jit, grad, vmap, and jvp. Numerical compatibility does not by itself guarantee ETP compiler recognition: a transformation may decompose a marker primitive into ordinary JAX operations. When a transformed function will be compiled for online learning, inspect the compiled graph and prefer the direct batched ETP operator where necessary.

x = jnp.ones((4, 3))
w = jnp.ones((3, 5))

# ---- JIT compilation ----
y_jit = jax.jit(braintrace.matmul)(x, w)
print("JIT output shape:", y_jit.shape)

JIT output shape: (4, 5)

# ---- Gradient computation ----
grad_fn = jax.grad(lambda w: jnp.sum(braintrace.matmul(x, w)))
dw = grad_fn(w)
print("Gradient shape:", dw.shape)
print("Gradient values:\n", dw)

Gradient shape: (3, 5)
Gradient values:
 [[4. 4. 4. 4. 4.]
 [4. 4. 4. 4. 4.]
 [4. 4. 4. 4. 4.]]

# ---- Vectorized mapping (vmap) ----
# vmap over a batch of inputs, each of shape (4, 3)
xs = jnp.ones((8, 4, 3))  # 8 different batches
vmap_fn = jax.vmap(lambda x_i: braintrace.matmul(x_i, w))
ys = vmap_fn(xs)
print("vmap output shape:", ys.shape)  # (8, 4, 5)

vmap output shape: (8, 4, 5)

# ---- JVP (forward-mode differentiation) ----
primals = (x, w)
tangents = (jnp.ones_like(x), jnp.ones_like(w))

y_primal, y_tangent = jax.jvp(braintrace.matmul, primals, tangents)
print("JVP primal shape:", y_primal.shape)
print("JVP tangent shape:", y_tangent.shape)

JVP primal shape: (4, 5)
JVP tangent shape: (4, 5)

# ---- Composability: JIT + grad + vmap ----
@jax.jit
def batched_grad(xs, w):
    """Compute per-sample gradients w.r.t. the weight."""
    def single_grad(x_i):
        return jax.grad(lambda w_: jnp.sum(braintrace.matmul(x_i, w_)))(w)
    return jax.vmap(single_grad)(xs)

xs = jnp.ones((8, 4, 3))
per_sample_grads = batched_grad(xs, w)
print("Per-sample gradients shape:", per_sample_grads.shape)  # (8, 3, 5)

Per-sample gradients shape: (8, 3, 5)

Next steps#

Continue with Neural Network Layers for Online Learning to compose these operators into reusable models, then read Hidden States for Online Learning to make those models temporal. Extension authors can continue to Creating Custom ETP Primitives.