flax

b4d1895d - Add comment to prefetch_to_device explaining it is mostly useful for GPUs (not TPUs).

Add comment to prefetch_to_device explaining it is mostly useful for GPUs (not TPUs). Following a user discussion: this is mostly useful for TPU, so it seems useful to add this as documentation.