pin_memory=True pytorch custom dataloader support

[externalsupplierstaff]

Hi! I love jax because it's very transparent. However, when implementing a custom Pytorch Dataset and the subsequent dataloader I have the question of where is the data loaded to. More specifically, in the docs' tutorial about dataloading, there is no device, cuda or gpu. I also see in the documentation that the (pre)allocation is automagically done.

My questions are the following:

• is the preallocation similar or equivalent than pin_memory=True?
• to put my variables in the GPU it is enough to convert them to jnp.float32 type, as long as gpu is available, true or false?
• why isn't this more transparent like pytorch's to(device)?

[externalsupplierstaff]

In terms of devices, what is the difference between the following? What does the copy_to_device function do?

import numpy as np
import jax.numpy as jnp
arrays = [np.array([3.14]),
jnp.array([3.14]),
np.array([3.14], dtype=jnp.float32),
jnp.array([3.14], dtype=jnp.float32)]
for arr in arrays:
  print(arr)
  print(type(arr))
  print(arr.dtype)
  print('-'*50)