Compatibility for 2.7 release

[zpcore]

jax_env_context moved its location. Make the update to make sure 2.7 release is runnable.

[tengyifei]

I think the old import works fine. Does it stop working in 2.8.0 for some reason? The versioning check is extra complexity.

[tengyifei]

Oh, never mind. The current code doesn't work on 2.7.

In order to make it work on both 2.7 and 2.8, I think simply

from torch_xla.experimental.custom_kernel import jax_env_context

Should work. LMK if not the case.

[zpcore]

Unfortunately, jax_env_context is not under path torch_xla.experimental.custom_kernel anymore in the PyTorch/XLA head.

[tengyifei]

ugh. you're right

[tengyifei]

Later on, it would be good to add an E2E test verifying certain recipes with the 2.7 docker. Otherwise, we'd easily regress and random things may start breaking as changes get landed.

[zpcore]

This is not enough, there is a API compatibility issue which you fixed before:

File "/workspaces/torchprime/torchprime/torch_xla_models/remat_all.py", line 29, in remat_all_partition_fn
    return min_cut_rematerialization_partition(
TypeError: min_cut_rematerialization_partition() got an unexpected keyword argument 'static_lifetime_input_indices'"

I guess we can drop this PR and focusing on syncing docker image periodically.

[tengyifei]

Maybe we could add an install sh script that downloads the right nightly version wheel corresponding to the pinned docker and installs them. To make the local installation process easier.