Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support nogil mode in JAX for Python 3.13 #23073

Open
1 task
vfdev-5 opened this issue Aug 15, 2024 · 2 comments
Open
1 task

Add support nogil mode in JAX for Python 3.13 #23073

vfdev-5 opened this issue Aug 15, 2024 · 2 comments
Assignees
Labels
enhancement New feature or request

Comments

@vfdev-5
Copy link
Collaborator

vfdev-5 commented Aug 15, 2024

The idea is to add a support of nogil mode in JAX for Python 3.13. According to Jake and Peter, this would likely involve updates to both jaxlib and to nanobind.

Context:

cc @jakevdp, @hawkinsp

@vfdev-5 vfdev-5 added the enhancement New feature or request label Aug 15, 2024
@vfdev-5 vfdev-5 self-assigned this Aug 15, 2024
vfdev-5 added a commit to vfdev-5/jax that referenced this issue Aug 20, 2024
… and `_tpu_ext`

Description:
- Added `py::mod_gil_not_used()` to `PYBIND11_MODULE` for `_triton_ext` and `_tpu_ext`.

Refs:
- https://py-free-threading.github.io/porting/#__tabbed_1_2

Context:
- jax-ml#23073
nitins17 pushed a commit to google-ml-infra/jax-fork that referenced this issue Aug 27, 2024
… and `_tpu_ext`

Description:
- Added `py::mod_gil_not_used()` to `PYBIND11_MODULE` for `_triton_ext` and `_tpu_ext`.

Refs:
- https://py-free-threading.github.io/porting/#__tabbed_1_2

Context:
- jax-ml#23073
@wjakob
Copy link

wjakob commented Sep 20, 2024

Nanobind's master branch now supports free-threading. Here is more information about the feature: https://nanobind.readthedocs.io/en/latest/free_threaded.html

@hawkinsp
Copy link
Collaborator

Thanks @wjakob . We're loving nanobind, by the way!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

No branches or pull requests

3 participants