jax
fdb1c144 - Switches make_array_from_callback to use batched_device_put

Commit
1 year ago
Switches make_array_from_callback to use batched_device_put PiperOrigin-RevId: 647537267
Author
Committer
Parents
Loading