You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
I'm trying to clarify the behavior of jax.experimental.io_callback when used with the ordered=True argument.
When ordered=True is set, does JAX guarantee that multiple io_callback invocations will be executed sequentially on the host, i.e., the next callback only starts after the previous one finishes?
Or, does JAX still launch multiple threads for the callbacks, and it is still necessary to manually add locking inside the callback function to prevent race conditions, e.g., when accessing global variables, writing to files, or using libraries like TensorBoardX?
reacted with thumbs up emoji reacted with thumbs down emoji reacted with laugh emoji reacted with hooray emoji reacted with confused emoji reacted with heart emoji reacted with rocket emoji reacted with eyes emoji
Uh oh!
There was an error while loading. Please reload this page.
-
Hi all,
I'm trying to clarify the behavior of
jax.experimental.io_callback
when used with theordered=True
argument.When
ordered=True
is set, does JAX guarantee that multipleio_callback
invocations will be executed sequentially on the host, i.e., the next callback only starts after the previous one finishes?Or, does JAX still launch multiple threads for the callbacks, and it is still necessary to manually add locking inside the callback function to prevent race conditions, e.g., when accessing global variables, writing to files, or using libraries like TensorBoardX?
Beta Was this translation helpful? Give feedback.
All reactions