In [9]: with jax.default_device(jax.devices("cpu")[0]):
   ...:     a = torch.randn((2,2), device='jax')
   ...:

In [10]: a._elem.device
Out[10]: CpuDevice(id=0)

In [11]: with jax.default_device(jax.devices("cpu")[0]):
    ...:     a = torch.randn((2,2)).to(device='jax')
    ...:

In [12]: a
Out[12]:
Tensor(<class 'jaxlib._jax.ArrayImpl'> [[0.7494194  0.94113636]
 [0.18599068 0.8405661 ]])

In [13]: a._elem.device
Out[13]: CpuDevice(id=0)

In [14]: with jax.default_device(jax.devices("tpu")[0]):
    ...:     a = torch.randn((2,2)).to(device='jax')
    ...:

In [15]: a._elem.device
Out[15]: TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0)