brainevent.jitumm

Contents

brainevent.jitumm#

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

Float matrix-matrix product with a JIT uniform connectivity matrix.

Computes the product of a just-in-time generated sparse matrix with uniformly distributed weights and a dense matrix B. Unlike the binary variant, this function uses the full floating-point values of B.

The sparse matrix A of shape (m, n) is never materialized. Each entry A[i, j] is drawn from Uniform(w_low, w_high) with probability prob, seeded by seed.

Parameters:

  • w_low (Array | ndarray | Quantity | Number) – Lower bound of the uniform weight distribution. Scalar value, optionally with physical units (brainunit.Quantity).

  • w_high (Array | ndarray | Quantity | Number) – Upper bound of the uniform weight distribution. Must have the same dimension (units) as w_low.

  • prob (float) – Connection probability in [0, 1]. Determines the fraction of non-zero entries in the connectivity matrix.

  • B (Array | ndarray | Quantity | Number) – Input dense matrix of shape (n, k) (if transpose=False) or (m, k) (if transpose=True). Optionally with physical units.

  • seed (int | None) – Random seed for reproducible connectivity patterns. If None, a random seed is generated at compile time.

  • shape (Tuple[int, int]) – Shape (m, n) of the logical connectivity matrix.

  • transpose (bool) – If True, compute A.T @ B instead of A @ B. Default is False.

  • corder (bool) – Memory layout order for the connectivity generation. True for C-order (row-major), False for Fortran-order (column-major). Default is True.

  • backend (str | None) – Computation backend. One of 'numba' or 'pallas'. If None, the default backend is used.

Returns:

Result matrix of shape (m, k) (if transpose=False) or (n, k) (if transpose=True). Carries the product of units from the weight and B if either has physical units.

Return type:

Array | ndarray | Quantity | Number

See also

jitumv

Matrix-vector variant.

binary_jitumm

Event-driven (binary) variant.

Notes

The connectivity matrix A of shape (m, n) follows the model:

A[i, j] = U[i, j] * B_conn[i, j]

where U[i, j] ~ Uniform(w_low, w_high) and B_conn[i, j] ~ Bernoulli(prob) are independent, both determined by seed.

The float matrix-matrix product computes:

result[i, j] = sum_{k=0}^{n-1} A[i, k] * B[k, j]

Unlike the binary variant (binary_jitumm()), this uses the full floating-point values of B rather than treating them as binary events.

When transpose=True, the operation becomes result = A^T @ B:

result[j, l] = sum_{i=0}^{m-1} A[i, j] * B[i, l]

Examples

>>> import jax.numpy as jnp
>>> from brainevent._jit_uniform.float import jitumm
>>> B = jnp.ones((5, 3))
>>> result = jitumm(0.1, 0.5, 0.2, B, seed=42, shape=(4, 5))
>>> result.shape
(4, 3)
Contents