brainevent.csrmm_dt2t

Contents

brainevent.csrmm_dt2t#

brainevent.csrmm_dt2t = <NameScope(brainevent.csrmm_dt2t)>#

Batched element-wise product of dense vectors and CSR weights, indexed by CSR structure.

Batched generalization of csrmv_dt2t(): for each batch element b and non-zero entry j in the CSR matrix at position (row, col), computes out[b, j] = w[b, j] * y[b, row] (non-transposed) or out[b, j] = w[b, j] * y[b, col] (transposed). Both w and y carry a shared leading batch axis, and the result has the same shape as w (one value per batch element and per structural non-zero).

This operation implements the \(D^t \epsilon^{t-1}\) term of the D-RTRL eligibility-trace update \(\epsilon^t \approx D^t \epsilon^{t-1} + \mathrm{diag}(D_f^t) \otimes x^t\) for a batch of samples: y holds the batched per-neuron factor \(D^t\) and w holds the batched per-synapse trace \(\epsilon^{t-1}\) stored in CSR data order.

The function supports physical units via brainunit.

Parameters:

  • y (Array | ndarray | Quantity | Number) – Batch of dense vectors indexed by the CSR structure. Shape (n_batch, shape[0]) when transpose=False or (n_batch, shape[1]) when transpose=True.

  • w (Array | ndarray | Quantity | Number) – Batch of per-synapse values. Shape (n_batch, nse); the trailing axis must match the shape of indices.

  • indices (Array | ndarray) – Column indices of the CSR matrix. Shape (nse,) with integer dtype.

  • indptr (Array | ndarray) – Row index pointer array. Shape (shape[0] + 1,) with integer dtype.

  • shape (tuple of int) – Two-element tuple (m, k) giving the logical shape of the CSR matrix.

  • transpose (bool) – If True, index y by column indices instead of row indices. Default is False.

  • backend (str | None) – Compute backend. Default is None (auto-select).

Returns:

out – Batched per-synapse result. Shape (n_batch, nse), same as w.

Return type:

Array | ndarray | Quantity | Number

See also

csrmv_dt2t

Unbatched (single-vector) variant of this operator.

csrmm

Standard CSR matrix–matrix multiplication.

Notes

This operation is differentiable with respect to both y and w via custom JVP rules. The transpose rule is not yet implemented.

Mathematically, for each batch element b and structural non-zero entry j of the CSR matrix at position (row, col), the output is computed as:

out[b, j] = w[b, j] * y[b, row] (non-transposed, transpose=False)

out[b, j] = w[b, j] * y[b, col] (transposed, transpose=True)

where row is determined by the indptr array (the row to which the j-th non-zero belongs) and col = indices[j].

Like csrmv_dt2t(), the output keeps the unit of w and drops the unit of y.

Examples

>>> import jax.numpy as jnp
>>> from brainevent import csrmm_dt2t
>>> y = jnp.array([[1.0, 2.0], [10.0, 20.0]])           # (n_batch=2, m=2)
>>> w = jnp.array([[0.5, 0.3, 0.7, 0.1],
...                [1.0, 2.0, 3.0, 4.0]])               # (n_batch=2, nse=4)
>>> indices = jnp.array([0, 2, 1, 2], dtype=jnp.int32)
>>> indptr = jnp.array([0, 2, 4], dtype=jnp.int32)
>>> csrmm_dt2t(y, w, indices, indptr, shape=(2, 3))
Contents