brainevent.binary_densemv_p

Contents

brainevent.binary_densemv_p#

brainevent.binary_densemv_p = <brainevent.XLACustomKernel object>#

Low-level XLA custom-kernel primitive for binary_densemv.

This XLACustomKernel instance dispatches the binary (event-driven) dense matrix-vector multiplication operation to registered backends. On GPU, cuda_raw is the WPR event-driven CUDA backend and remains the default; jax_raw provides a pure JAX reference path, and cublas provides an optional float32/bool cuBLAS path.

The operation computes weights[m,k] @ spikes[k] -> out[m] when transpose=False, or spikes[k] @ weights[k,n] -> out[n] when transpose=True, where only active (nonzero) spike entries contribute to the output.

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

See also

binary_densemv

High-level user-facing function wrapper.