jax
8699f5d9 - When host local inputs on all hosts are the same, use `_DeferredShardArg` to do the transfers instead of `jit` to avoid blocking.

Commit

1 year ago

When host local inputs on all hosts are the same, use `_DeferredShardArg` to do the transfers instead of `jit` to avoid blocking. PiperOrigin-RevId: 699336402

References

#25074 - When host local inputs on all hosts are the same, use `_DeferredShardArg` to do the transfers instead of `jit` to avoid blocking.

#170 - CI: 12/02/24 upstream sync

Author

yashk2810

Committer

Google-ML-Automation

Parents

030ee4a1

jax 8699f5d9 - When host local inputs on all hosts are the same, use `_DeferredShardArg` to do the transfers instead of `jit` to avoid blocking.

jax
8699f5d9 - When host local inputs on all hosts are the same, use `_DeferredShardArg` to do the transfers instead of `jit` to avoid blocking.