Hello all,
is there some way to load a JAX array into a torch tensor? A naive way of doing this would be
import numpy as np

np_array = np.asarray(jax_array)
torch_ten = torch.from_numpy(np_array).cuda()

This would be slow as it would require me to move the jax array from the gpu to a cpu numpy array before loading it on the gpu again.
Just to be clear: I am not interested in any gradient information being preserved. Only the entries of the array need to be loaded.
Many thanks!