brainevent.binary_coomm_p

Contents

brainevent.binary_coomm_p#

brainevent.binary_coomm_p = <brainevent.XLACustomKernel object>#

Low-level XLA custom-kernel primitive for binary_coomm.

This XLACustomKernel instance dispatches the binary (event-driven) COO sparse matrix-matrix multiplication operation to registered backends (numba, pallas), using runtime shape/dtype metadata provided by the high-level wrapper.

The operation computes result[i, n] = sum_j A[i, j] * (B[j, n] > 0) when transpose=False or result[j, n] = sum_i A[i, j] * (B[i, n] > 0) when transpose=True, where only active (nonzero) events in the dense matrix B contribute to the output. This is efficient for processing batches of spike events in neural network simulations.

Beyond backend dispatch, the primitive stores JAX transformation bindings (JVP, transpose, batching, and call registration) so the operation integrates correctly with jit, vmap, and autodiff.

Available backends can be queried with binary_coomm_p.available_backends(platform), and the default backend can be configured with binary_coomm_p.set_default(platform, backend).

See also

binary_coomm

High-level user-facing function wrapper.