jax

3146c2a3 - Merge pull request #20725 from mattjj:io-callback-batching-fix