$ the-wire · showcase
JAX refactors array construction for free-threading safety
By RepoJournal · Filed · About Google
JAX's array construction is being refactored to prepare for Python free threading, removing unsafe internals and moving validation out of the constructor.
JAX first [1] relocates shard validation and rearrangement from the PyArray C++ constructor into PyArray::MakeFromIfrtArrayAndSharding, converting _check_and_rearrange in array.py to a staticmethod that takes (arrays, sharding, aval), and removes skip_checks from the constructor. To support concurrency under Python free threading, [2] adds PyArray::ifrt_array_ref(), which returns xla::ifrt::ArrayRef instead of a raw pointer; the old ifrt_array() method stays temporarily for backwards compatibility. [3] deletes the skip_checks parameter from PyArrayResultHandler, since it was always called with skip_checks=true, and guards Python call sites via jaxlib_extension_version >= 489. [4] changes pallas:collective_ids' auto-ID assignment to scan the Jaxpr during lowering for manually assigned IDs and allocate non-colliding ones instead of using a base offset. Separately in googleapis/google-cloud-python, [5] onboarded the google-cloud-apptopology library using Librarian, but [6] reverted it because a checked-in scripts directory shouldn't have been included, then [7] re-landed the same feature. Spanner test flakes are reduced: [8] fixes test_create_instance_partition by using a multi-region configuration to match a Cloud Spanner validation rule.
Action items
- → Check whether you rely on PyArray's ifrt_array() method; it is being removed in a follow-up after the new ifrt_array_ref() lands. google/jax [plan]
- → Verify your prebuilt jaxlib version is at least 489 if you run custom Python call sites, as older versions still require _skip_checks=True. google/jax [immediate]
- → Review the re-landed google-cloud-apptopology onboarding in google-cloud-python to ensure no unintended files are included. googleapis/google-cloud-python [monitor]
References
- [1] [JAX] Remove _check_and_rearrange and skip_checks from PyArray constructor. ↗ google/jax
- [2] Add `PyArray::ifrt_array_ref()`, which returns `xla::ifrt::ArrayRef` instead of a raw pointer. ↗ google/jax
- [3] Remove the skip_checks parameter from PyArrayResultHandler. ↗ google/jax
- [4] [pallas:collective_ids] Remove base offset in favor of non-colliding auto IDs. ↗ google/jax
- [5] feat(google/cloud/apptopology/v1): add google-cloud-apptopology ↗ googleapis/google-cloud-python
- [6] Revert "feat(google/cloud/apptopology/v1): add google-cloud-apptopology" ↗ googleapis/google-cloud-python
- [7] feat(google/cloud/apptopology/v1): add google-cloud-apptopology ↗ googleapis/google-cloud-python
- [8] chore(spanner): fix instance partition test configuration and snippets test flakes ↗ googleapis/google-cloud-python
