Sharding Support in XLA Compiler

[ntenenz]

@mattjj recently presented at GTC a cleaner API for working with ShardedDeviceArrays that leans more heavily on the XLA compiler to produce parallelized versions of algorithms. This appears to be undocumented in the python API at this time (e.g., device_put), leading to two inevitable questions...

• Is this support available in nightly today?
• If not, is there an ETA as to when this will be released?

[yashk2810]

Hello, yeah it's undocumented currently as we are working on rolling it out.

Is this support available in nightly today?

Yes, you can install jax at HEAD, jaxlib-nightly and/or libtpu-nightly and get the features that were presented. Don't forget to enable jax_array flag!

If not, is there an ETA as to when this will be released?

We are working on enabling this functionality globally and will take some time. I can't give you an exact date sadly, but you can try it out!

[ntenenz]

Happy to work from HEAD. Thanks for the confirmation. :)