mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-21 22:36:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
26,443 Commits 508 Branches 263 Tags
rocm-main
Commit Graph

2 Commits

Author SHA1 Message Date
Peter Hawkins
64eae324ee Migrate JAX MLIR Python dialect extensions to nanobind. 
Now that https://github.com/llvm/llvm-project/pull/117922 has landed upstream, we can work towards removing our last uses of pybind11.

PiperOrigin-RevId: 705872751
 2024-12-13 07:08:28 -08:00
Peter Hawkins
70b7d50181 Switch jaxlib to use nanobind instead of pybind11. 
nanobind has a number of advantages (https://nanobind.readthedocs.io/en/latest/why.html), notably speed of compilation and dispatch, but the main reason to do this for these bindings is because nanobind can target the Python Stable ABI starting with Python 3.12. This means that we will not need to ship per-Python version CUDA plugins starting with Python 3.12.

PiperOrigin-RevId: 559898790
 2023-08-24 16:07:56 -07:00