Synchronization

Host/device synchronization helpers for blocking on and copying pytrees of JAX arrays.

See Also

• Performance Overview - All performance tools
• Goodput - Goodput tracking
• Roofline - Roofline analysis

---

datarax.performance.synchronization

Shared synchronization helpers for timing JAX work.

block_until_ready_tree

block_until_ready_tree(value: Any) -> Any

Synchronize every JAX-like leaf in a pytree and return the original value.

copy_to_host_async_tree

copy_to_host_async_tree(value: Any) -> Any

Start async host copies for every supporting leaf and return the original value.