brainevent.coomv_p

brainevent.coomv_p = <brainevent.XLACustomKernel object>

Low-level XLA custom-kernel primitive for coomv.

This XLACustomKernel instance dispatches the COO sparse matrix-vector multiplication with floating-point weights operation to registered backends (numba, pallas), using runtime shape/dtype metadata provided by the high-level wrapper.

The operation computes y[i] = sum_k A[i, k] * v[k] when transpose=False or y[k] = sum_i A[i, k] * v[i] when transpose=True, where A is a sparse matrix defined by (data, row, col) in COO format and v is a dense vector. This is the standard sparse matrix-vector multiplication with full floating-point arithmetic.

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 coomv_p.available_backends(platform), and the default backend can be configured with coomv_p.set_default(platform, backend).

See also

coomv

High-level user-facing function wrapper.