Tries to call a copy_to_host_async method on pytree leaves.
For each leaf this method will try to call the copy_to_host_async method on the leaf. If the leaf is not a JAX array, or if the leaf does not have a copy_to_host_async method, then this method will do nothing to the leaf.
x – a pytree, usually with at least some JAX array instances at its leaves.
A pytree with the same structure and values of the input, where the host copy of the values of all JAX array leaves are started.
RetroSearch is an open source project built by @garambo | Open a GitHub Issue
Search and Browse the WWW like it's 1997 | Search results from DuckDuckGo
HTML:
3.2
| Encoding:
UTF-8
| Version:
0.7.3