jax

5e915d33 - Update the sparse GPU kernels in jaxlib to use the FFI.

Update the sparse GPU kernels in jaxlib to use the FFI. Unlike the other more detailed ports, this version doesn't take full advantage of the features provided by the FFI. For example, it would be possible to update the kernels to use the ScratchAllocator instead of querying the workspace size during lowering. However, since these kernels are really only meant to be experimental, it's not obvious to me that it's worth the extra work to do anything more sophisticated. PiperOrigin-RevId: 724016331