brainevent.csrmm_dt2t_p

Contents

brainevent.csrmm_dt2t_p#

brainevent.csrmm_dt2t_p = <brainevent.XLACustomKernel object>#

Low-level XLA custom-kernel primitive for csrmm_dt2t.

This XLACustomKernel instance 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 with csrmm_dt2t_p.set_default(platform, backend).

See also

csrmm_dt2t

High-level user-facing function wrapper.