Ship torchax inside torch_xla

[tengyifei]

🚀 Feature

When users install torch_xla, they would get torchax as well.

Motivation

The xb.call_jax 1 feature under heavy development uses torchax and we expect torchax usage to grow. In order for all users of torch_xla to use this feature, we need to include torchax in torch_xla wheels/dockers/etc.

Pitch

I think there are two approaches:

1. Make torchax its own Python wheel/package, and have torch_xla depend on it always.

2. Literally link all torchax files inside the torch_xla wheel/package. This approach might need to deal with some packaging subtleties because some users may have already installed torchax standalone and not sure what will happen when they also install torch_xla with its own torchax.

Either way, this probably requires some setup.py changes.

cc @qihqi

[tengyifei]

A big motivation is making atomic changes across torchax and torch_xla. This is only doable if torch_xla includes torchax by source. So the wheel option is probably out.