Is there an interface in jax similar to copyFromHostToDevice in TensorFlow's Swap

[Mewo518]

I'd like to implement the swap and recompute functions in memory optimization in jax, where is this interface function

[jakevdp]

I don't know what copyFromHostToDevice is, but perhaps jax.device_put is what you have in mind? If not, you might edit your question to add more detail about exactly what you're looking for.

Hope that helps!