Add an e2e test for environment setup

[vlasenkoalexey]

Currently torchprime is built from torch_xla head, that in turns pulls in pinned version of libtpu.
This setup is fragile and could break easily, we need a test to confirm that it works to reduce friction for new users.

[tengyifei]

My idea:

1. Add a ./install_torch_xla.sh that does this

pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
# TODO: it should replace `cp310-cp310` with the right Python version
pip install 'torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.8.0.dev-cp310-cp310-linux_x86_64.whl' \
  -f https://storage.googleapis.com/libtpu-wheels/index.html

2. Test this in a CPU workflow/CI

3. Document on the readme that people can just call ./install_torch_xla.sh to get the nightly

[jialei777]

is that possible to not using nightly version?

[tengyifei]

is that possible to not using nightly version?

At least right now it's hard. See #213.

One complication is that let's say we want to move a feature (e.g. host offloading from torchprime to torch_xla). Then that requires bumping the torch_xla dep to include that upstreamed feature.

[yaoshiang]

One line in the sand could be to say, we only work on latest release (e.g. 2.7)

A slightly relaxed statement would be, we work on latest plus some patches that are contained in torchprime itself. Over time, those patches get upstreamed into PT/XLA. This basically allows us to patch any python based PT/XLA.

An even more relaxed statement is that we work on PT/XLA nightly. It can be reasonable in general for well tested nightlies, though I am not certain of the ci/cd quality of PT/XLA nightly (I think torchbench tests of pt/xla failed for 30 days straight with no fixes).

The most relaxed statement of building from head, I would say does not serve the vision of torchprime:

a) an easy-to-learn-from reference implementation (PT/XLA features in head are inherently unstable, subject to change, poorly documented, etc)
b) an easy-to-use tool (building PT/XLA from head is not easy)
c) showing off PT/XLA performance (statement from Dylan Patel: "We believe that the correct way to interpret the results is to look at the performance of the public stable release of [xxx] software.")