brainevent.csrmv_yw2y_p

Contents

brainevent.csrmv_yw2y_p#

brainevent.csrmv_yw2y_p = <brainevent.XLACustomKernel object>#

Low-level XLA custom-kernel primitive for csrmv_yw2y.

This XLACustomKernel instance dispatches the CSR sparse matrix-vector multiplication with element-wise product (y * W -> y) operation to registered backends (numba, pallas), using runtime shape/dtype metadata provided by the high-level wrapper.

For each non-zero entry at position (row, col) in the CSR matrix, computes out[j] = w[j] * y[row] (non-transposed) or out[j] = w[j] * y[col] (transposed), producing one output element 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 csrmv_yw2y_p.available_backends(platform), and the default backend can be configured with csrmv_yw2y_p.set_default(platform, backend).

See also

csrmv_yw2y

High-level user-facing function wrapper.