$ the-wire · showcase
JAX ships thread-safe PyArray, changes jax.jit inline semantics
By RepoJournal · Filed · About Google
JAX merged three changes today: PyArray is now thread-safe, inline=True in jax.jit means late inlining, and a previous commit was reverted.
JAX made PyArray thread-safe by adding a mutex (`storage.mu`) to each array, guarding its buffer, cached NumPy value, shards, and status. Accessing a deleted array now returns an error instead of crashing, and replacing a buffer (`_replace_with`) while another thread reads or shards it raises an error. Host NumPy arrays are marked read-only before publishing, so other threads can't mutate them behind your back. This is a deep concurrency hardening for anyone sharing arrays across threads [1].
JAX changed `inline=True` in `jax.jit` to mean `jax.Inline.JAX_LATE` (inlining during MLIR lowering) instead of `jax.Inline.JAX_EARLY` (during Python tracing into jaxpr). The effect: calls stay in jaxpr representations but get inlined before XLA compilation. Existing uses inside JAX itself were updated to `inline=JAX_EARLY` to preserve their behavior, but external code using `inline=True` will now take the late path unless explicitly changed [2].
RepoJournal also reverted a JAX commit, and googleapis cleaned up python-genai's FunctionResultDelta by removing an unused field [3]. For bigtable in google-cloud-python, `MutateRows` now routes through the data client's `Table.bulk_mutate_rows`, reworking error handling to rebuild a list of `Status` objects from the `_MutateRowsOperation` errors [4]. The storage README fix resolves a reST target mismatch that caused PyPI upload failures and sets `long_description_content_type` explicitly [5].
Action items
- → Review jax.jit uses with inline=True if you rely on eager inlining during tracing google/jax [monitor]
- → Test concurrent access to JAX arrays if you share them across threads google/jax [plan]
References
- [1] Make PyArray thread-safe. ↗ google/jax
- [2] Change inline=True in jax.jit to mean JAX_LATE instead of JAX_EARLY. ↗ google/jax
- [3] fix: Remove unused field from FunctionResultDelta ↗ googleapis/python-genai
- [4] feat(bigtable): Rerouted MutateRows to use the data client ↗ googleapis/google-cloud-python
- [5] fix(storage): resolve reST link target mismatch in README and add long_description_content_type (#18210) ↗ googleapis/google-cloud-python
