jax
c69f2ec7 - [Roll forward -- Take 3] Optimize large JAX Mesh construction

Commit
41 days ago
[Roll forward -- Take 3] 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 d925891ab3270b1273854b47b60f1ad5732579cf PiperOrigin-RevId: 983872140
Author
Parents
Loading