Dan Foreman-Mackey b308c64936 Export jaxlib.xla_client.register_custom_call_target as jax.extend.ffi.register_ffi_target. ...

This means that users of the FFI interface won't need to directly
interact with `jaxlib.xla_client` at all.

I've expanded the doctring a little and changed one default: the default
`api_version` is `1` instead of `0` to be consistent with the new name.

2024-07-19 08:12:25 -04:00

..

__init__.py

add jax.extend.random.wrap_key_data

2023-08-26 11:39:25 -07:00

ffi.py

Export jaxlib.xla_client.register_custom_call_target as jax.extend.ffi.register_ffi_target.

2024-07-19 08:12:25 -04:00

random.py

Run pyupgrade --py310-plus.

2024-06-26 16:10:18 -04:00