brainevent.binary_coomm_p#
- brainevent.binary_coomm_p = <brainevent.XLACustomKernel object>#
Low-level XLA custom-kernel primitive for
binary_coomm.This
XLACustomKernelinstance 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)whentranspose=Falseorresult[j, n] = sum_i A[i, j] * (B[i, n] > 0)whentranspose=True, where only active (nonzero) events in the dense matrixBcontribute 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 withbinary_coomm_p.set_default(platform, backend).See also
binary_coommHigh-level user-facing function wrapper.