brainevent.fcnmm_dt2t

Contents

brainevent.fcnmm_dt2t#

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

Batched per-synapse element-wise product of dense vectors and fixed-connection weights.

Batched generalization of fcnmv_dt2t(): for each batch element b and stored connection slot (i, k) of the fixed-connection-number (ELL) structure, computes out[b, i, k] = weights[b, i, k] * y[b, i] when transpose=False or out[b, i, k] = weights[b, i, k] * y[b, indices[i, k]] when transpose=True. Both weights and y carry a shared leading batch axis, and the result has shape (n_batch, *indices.shape).

This is the fixed-connection-number analog of brainevent.csrmm_dt2t() and follows the same transpose convention as that operator: the flag selects whether y is indexed by the leading (row) axis or by the stored indices (column) axis. Note this is the opposite sense of the transpose flag in brainevent.fcnmv() / brainevent.fcnmm().

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 weights holds the batched per-synapse trace \(\epsilon^{t-1}\) stored in the ELL grid.

The operator is a plain element-wise / gather-multiply, so it is fully differentiable through JAX autodiff with respect to both weights and y; no custom differentiation rule is required. It is also jit- and vmap-compatible.

Parameters:

  • weights (Array | Quantity) – Batched per-synapse values. Either a scalar / size-1 array (homogeneous) or a (n_batch, rows, n_conn) array whose trailing axes match indices (heterogeneous). Must be a floating-point dtype.

  • indices (Array) – Integer ELL index array of shape (rows, n_conn).

  • y (Array | Quantity) – Batch of dense vectors. Shape (n_batch, shape[0]) when transpose=False or (n_batch, shape[1]) when transpose=True.

  • shape (Tuple[int, int]) – Two-element (rows_dim, cols_dim) logical shape of the ELL operand.

  • transpose (bool) – If False, index y by the leading (row) axis (broadcast over the connection axis). If True, index y by indices (gather).

Returns:

out – Batched per-synapse result of shape (n_batch, *indices.shape). The output unit is unit(weights) * unit(y).

Return type:

Array | Quantity

Raises:

  • ValueError – If indices is not 2-D.

  • ValueError – If shape is not length-2.

  • ValueError – If weights is not a floating-point dtype.

  • ValueError – If weights is neither size-1 nor shaped (n_batch, rows, n_conn).

  • ValueError – If y is not 2-D.

  • ValueError – If y does not have the trailing dimension required by transpose and shape.

See also

fcnmv_dt2t

Unbatched (single-vector) variant of this operator.

csrmm_dt2t

CSR equivalent of this operator.

fcnmm

Fixed-connection sparse matrix–matrix product.

Notes

Unlike brainevent.csrmm_dt2t() (which drops y’s physical unit), this operator keeps both units, consistent with fcnmv_dt2t() and with the mathematics of a product.

Mixed floating dtypes between weights and y are promoted following JAX’s standard promotion rules (e.g. float32 * float64 -> float64); no equal-dtype constraint is imposed.

Empty structures are supported: an indices array with rows == 0 or n_conn == 0, or a y with zero batch elements, returns an array of the corresponding (empty) shape.

This operation is memory-bandwidth-bound and maps to a fused XLA broadcast-multiply (transpose=False) or gather-multiply (transpose=True); it is therefore implemented in pure JAX rather than as a custom kernel.

Examples

>>> import jax.numpy as jnp
>>> from brainevent import fcnmm_dt2t
>>> # n_batch=2, rows=2, n_conn=2
>>> weights = jnp.array([[[0.5, 1.0], [1.5, 2.0]],
...                      [[1.0, 1.0], [1.0, 1.0]]])
>>> indices = jnp.array([[0, 1], [1, 2]])
>>> y = jnp.array([[10.0, 20.0], [100.0, 200.0]])  # (n_batch=2, rows=2)
>>> fcnmm_dt2t(weights, indices, y, shape=(2, 3), transpose=False)
Array([[[  5.,  10.],
        [ 30.,  40.]],

       [[100., 100.],
        [200., 200.]]], dtype=float32)

>>> # transpose=True gathers y by the stored column indices
>>> y_post = jnp.array([[1.0, 2.0, 3.0], [10.0, 20.0, 30.0]])
>>> fcnmm_dt2t(weights, indices, y_post, shape=(2, 3), transpose=True)
Array([[[ 0.5,  2. ],
        [ 3. ,  6. ]],

       [[10. , 20. ],
        [20. , 30. ]]], dtype=float32)
Contents