Question: Does XLA support auto-sharding for PyTorch on GPUs

[neeldani]

Hi community 👋

I want to run an experiment of XLA's auto-sharding (Alpa) on a GPU. However, I see that pytorch-xla does not support auto-sharding on GPUs yet (source)

I had the following questions:

1. Is there a way I can use the auto-sharding functionality present in the openxla/xla repo for running my PyTorch experiment across multiple GPUs? If yes, is there a document that I can refer to
2. If not PyTorch, can I use Jax or Tensorflow? XLA does a bunch of optimizations apart from auto-sharding so is there a document on how to just run the auto-sharding functionality for any of these frameworks (similar to just running alpa)?