jax
dfd7d17c - [JAX] Use iota_reshape_dims and iota_transpose_perm in pxla, which is more efficient than tile_assignment_devices.

Commit

1 year ago

[JAX] Use iota_reshape_dims and iota_transpose_perm in pxla, which is more efficient than tile_assignment_devices. HloSharding V1 -> HloSharding V2. PiperOrigin-RevId: 639210975

References

#21566 - [JAX] Use iota_reshape_dims and iota_transpose_perm in pxla, which is more efficient than tile_assignment_devices.

Author

ZixuanJiang

Committer

a-googler

Parents

1edb94ec

jax dfd7d17c - [JAX] Use iota_reshape_dims and iota_transpose_perm in pxla, which is more efficient than tile_assignment_devices.

jax
dfd7d17c - [JAX] Use iota_reshape_dims and iota_transpose_perm in pxla, which is more efficient than tile_assignment_devices.