Notebook async serialization

[sh0416]

I usually use notebook for development.

However, the jax array serialization requires async ops which raise an error during notebook execution.

Is there any simple trick to avoid this error?

[sh0416]

I mean this code.

https://github.com/google/jax/blob/main/jax/experimental/array_serialization/serialization.py