[torchax] jax_tensor.copy_(cpu_tensor) fails

[tengyifei]

Reproducer:

import torch
import torchax
torchax.enable_globally()

cpu_tensor = torch.tensor([1.0, 2.0])
jax_tensor = torch.tensor([2.0, 3.0], device='jax')
jax_tensor.copy_(cpu_tensor)


ERROR:
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/google/home/yifeit/torchprime/.venv/lib/python3.11/site-packages/torchax/tensor.py", line 257, in __torch_function__
    return func(*args, **(kwargs or {}))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/google/home/yifeit/torchprime/.venv/lib/python3.11/site-packages/torchax/tensor.py", line 275, in __torch_dispatch__
    return self.env.dispatch(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/google/home/yifeit/torchprime/.venv/lib/python3.11/site-packages/torchax/tensor.py", line 489, in dispatch
    res = op.func(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/google/home/yifeit/torchprime/.venv/lib/python3.11/site-packages/torchax/ops/jaten.py", line 132, in _aten_copy
    x._elem = y._elem.astype(x._elem.dtype)
              ^^^^^^^
AttributeError: 'Tensor' object has no attribute '_elem'

Per PyTorch docs (https://pytorch.org/docs/stable/generated/torch.Tensor.copy_.html), copying between tensors on different devices should be supported.

cc @qihqi