brainevent.jitsmm_p

brainevent.jitsmm_p = <brainevent.XLACustomKernel object>

Low-level XLA custom-kernel primitive for jitsmm.

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

In this operation, the connectivity matrix has all weights set to a single scalar value, and the input matrix contains floating-point values. Each column of the input is processed independently in a standard matrix-matrix product.

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

See also

jitsmm

High-level user-facing function wrapper.