[Torch_xla2] torch_xla2.default_env() guard doesn't enforce XLATensor2

[zpcore]

🐛 Bug

torch_xla2.default_env() doesn't guarantee XLATensor2 device. Need explicitly move the tensor to 'jax' to convert into XLATensor2.

To Reproduce

import torch
import torch_xla2

env = torch_xla2.default_env()

with env:
  inputs = torch.randn(1)
  print(type(inputs))   ---> output torch.Tensor instead of XLATensor2

Environment

PyTorch/XLA git commit f52e202e825651374e44dbb1c79fab3724be7e0e

[qihqi]

Please try either:

import torch
import torch_xla2

env = torch_xla2.default_env()

with env:
  inputs = torch.randn(1, device='jax')
  print(type(inputs))   ---> output torch.Tensor instead of XLATensor2

OR,

import torch
import torch_xla2

env = torch_xla2.default_env()
env.config.use_torch_native_for_cpu_tensor = False # this will use XLATensor even for CPU

with env:
  inputs = torch.randn(1)
  print(type(inputs))   ---> output torch.Tensor instead of XLATensor2