mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
26,443 Commits 501 Branches 263 Tags
Commit Graph

5 Commits

Author SHA1 Message Date
Dan Foreman-Mackey
5e915d3307 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
 2025-02-06 11:45:57 -08:00
Peter Hawkins
45b871950e Fix a number of minor problems in the ROCM build. 
Change in preparation for adding more presubmits for AMD ROCM.

PiperOrigin-RevId: 667766343
 2024-08-26 17:04:01 -07:00
Srinivas Vasudevan
7dfc8ff49d Add batching rules to jax.lax.linalg.tridiagonal_solve. 
PiperOrigin-RevId: 555700103
 2023-08-10 16:25:59 -07:00
Peter Hawkins
172a831219 Switch JAX to use the OpenXLA repository. 2023-03-13 18:38:26 +00:00
Peter Hawkins
a852710a09 Merge CUDA and ROCM kernel code in jaxlib. 
The code for both CUDA and ROCM is almost identical, so with a small shim library to handle the differences we can share almost everything.

PiperOrigin-RevId: 483666051
 2022-10-25 07:23:34 -07:00