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

References

Author

marksandler2

Committer

a-googler

Parents