jax
2f808e9d - Fix error in custom call registration for some FFI functions

Commit
1 year ago
Fix error in custom call registration for some FFI functions We are getting the following errors: ``` Duplicate FFI handler registration for cu_threefry2x32_ffi on a platform CUDA Duplicate FFI handler registration for cu_lu_pivots_to_permutation on a platform CUDA ``` It seems that with the ffi registration mechanism based on `XLA_FFI_REGISTER_HANDLER` it is not possible anymore to register a call target twice. The fix here is to rollback the changes in https://github.com/google/jax/pull/22178 and disable the changes from https://github.com/google/jax/pull/20997. PiperOrigin-RevId: 647993991
Author
Committer
Parents
Loading