brainevent.coomm

Contents

brainevent.coomm#

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

Perform COO sparse matrix-matrix multiplication.

Computes the product of a sparse matrix stored in COO format and a dense matrix B.

With transpose=False:

Y[i, n] = sum_{k} A[i, k] * B[k, n]

With transpose=True:

Y[k, n] = sum_{i} A[i, k] * B[i, n]

Parameters:

  • data (Array | ndarray | Quantity | Number) – Non-zero values of the sparse matrix. Either a scalar (shape (1,) for homogeneous weights) or a 1-D array of length nnz.

  • row (Array | ndarray) – 1-D int array of row indices, length nnz.

  • col (Array | ndarray) – 1-D int array of column indices, length nnz.

  • B (Array | ndarray | Quantity | Number) – Dense right-hand matrix. Shape (shape[1], n) when transpose=False, or (shape[0], n) when transpose=True.

  • shape (Tuple[int, int]) – Logical (m, k) shape of the sparse matrix.

  • transpose (bool) – If True, multiply by A.T. Default is False.

  • backend (str | None) – Compute backend. None selects the default.

Returns:

Result matrix. Shape (shape[0], n) when transpose=False, or (shape[1], n) when transpose=True.

Return type:

jax.Array or Quantity

See also

coomv

COO sparse matrix-vector multiplication.

binary_coomm

Event-driven (binary) COO matrix-matrix multiplication.

Notes

The kernel iterates over all nnz stored elements and, for each triplet (data[s], row[s], col[s]), adds data[s] * B[col[s], :] into Y[row[s], :] (forward) or data[s] * B[row[s], :] into Y[col[s], :] (transpose).

Physical units from brainunit are handled transparently.

Examples

>>> import jax.numpy as jnp
>>> from brainevent import coomm
>>> data = jnp.array([1.0, 2.0])
>>> row = jnp.array([0, 1])
>>> col = jnp.array([1, 0])
>>> B = jnp.ones((2, 3))
>>> coomm(data, row, col, B, shape=(2, 2))
Contents