xla

[torch_xla2] Add oss only tags to jax env initialization

#7920

Merged