rocm_jax

mirror of https://github.com/ROCm/jax.git synced 2025-04-14 10:56:06 +00:00

History

Daniel Suo 2d1bc5c2a0 Refactor Jax FFI lowering to prepare for implementing CPU/GPU callbacks using XLA's FFI.

- This refactor just moves code around and should have no impact on tests or public-facing APIs.
- `mlir.emit_python_callback` would eventually depend on `ffi.ffi_lowering`, which in turn depends on definitions in `mlir.py`. We break this circular dependency.

PiperOrigin-RevId: 729561359

2025-02-21 09:45:59 -08:00

__init__.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

ad.py

Remove a number of expired deprecations.

2024-10-31 15:40:54 -07:00

batching.py

Make trace dispatch purely a function of context rather than a function of both context and data. This lets us delete a lot of machinery for managing data-dependent tracing: levels, sublevels, post_process_call, new_base_main, custom_bind and so on.

2024-10-29 11:04:31 -07:00

mlir.py

Refactor Jax FFI lowering to prepare for implementing CPU/GPU callbacks using XLA's FFI.