google/jax
Pull Request Overview
- Opened on September 18, 2026
- Status Merged
- Commit count 1 with first commit September 18, 2026
Total Delta
0 Total Diff Delta
Open Days
Open 0 weekdays
Test Delta
0 Diff Delta in Test Files
Breakdown by Phase
How long has this pull request spent in each phase of its lifecycle?
Data pending calculation for pull request
[Roll forward -- Take 2] Optimize large JAX Mesh construction
[Roll forward -- Take 2] Optimize large JAX Mesh construction
- Avoid intermediate tuple -> np.ndarray reallocation: Previously Mesh._create took flat_devices_tuple and reallocated an np.ndarray via np.array(flat_devices_tuple).reshape(device_shape). We now introduce _HashableDeviceArray wrapping the original np.ndarray and xc.DeviceList, avoiding the array reallocation.
- Reuse xc.DeviceList: Cache _cached_device_list on the Mesh instance during construction and return it directly from the _internal_device_list property, avoiding redundant DeviceList allocations.
Reverts 008c954b905e7119157f55dc6ea7d16ee15ce05f
PR was closed without comments.