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
rocm_jax/jaxlib/rocm
History
Dan Foreman-Mackey 9ae1c56c44 Update lu_pivots_to_permutation to use FFI dimensions on GPU. 
The XLA FFI interface provides metadata about buffer dimensions, so quantities
like batch dimensions can be evaluated on the backend, instead of passed as
attributes. This change has the added benefit of allowing this FFI call to
support "vectorized" vmap and dynamic shapes.

PiperOrigin-RevId: 647343656
2024-06-27 09:27:15 -07:00
..
BUILD.bazel
Update lu_pivots_to_permutation to use FFI dimensions on GPU.
2024-06-27 09:27:15 -07:00