flax

dacbe35e - Add clarifying docstring for 'size' argument to prefetch_to_device's