jax

6bad8fcf - [pmap] Create new ifrt array when rewrapping global arrays as host local in multihost pmap.

[pmap] Create new ifrt array when rewrapping global arrays as host local in multihost pmap. Previously, the fast path `_rewrap_with_aval_and_sharding` we used when converting global arrays because to host-local arrays only updated the sharding on the PyArray and not the underlying ifrt array. As a result, supposedly host-local arrays returned from pmap would fail with the error below if users attempt to index (e.g., x[0]) in the multihost setting. `jax.errors.JaxRuntimeError: INVALID_ARGUMENT: CopyArrays only supports destination device list of the same size as the array device lists.` This change fixes that by also creating a new ifrt array using the new sharding and the disassembled single device arrays from the original ifrt array. We also add a multihost test that failed prior to this change and now succeeds. PiperOrigin-RevId: 895522456