brainevent.csrmv_yw2y_p#
- brainevent.csrmv_yw2y_p = <brainevent.XLACustomKernel object>#
Low-level XLA custom-kernel primitive for
csrmv_yw2y.This
XLACustomKernelinstance 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 withcsrmv_yw2y_p.set_default(platform, backend).See also
csrmv_yw2yHigh-level user-facing function wrapper.