brainevent.csrmm_dt2t_p#
- brainevent.csrmm_dt2t_p = <brainevent.XLACustomKernel object>#
Low-level XLA custom-kernel primitive for
csrmm_dt2t.This
XLACustomKernelinstance dispatches the batched CSR per-synapse element-wise product (y * w -> w-shaped output) operation to registered backends (numba,cuda_raw,jax_raw), using runtime shape/dtype metadata provided by the high-level wrapper.For each batch element b and non-zero entry at position (row, col) in the CSR matrix, computes out[b, j] = w[b, j] * y[b, row] (non-transposed) or out[b, j] = w[b, j] * y[b, col] (transposed), producing one output element per batch element and per structural non-zero.
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
csrmm_dt2t_p.available_backends(platform), and the default backend can be configured withcsrmm_dt2t_p.set_default(platform, backend).See also
csrmm_dt2tHigh-level user-facing function wrapper.