ResourceEnv requires 'positional argument: 'loops''

[mosmos6]

Hello
I'm on colab with JAX 0.3.5.
I'm trying to obtain arrays for TPU resource axis by ResourceEnv of jax.experimental.maps.
It works well by JAX0.2.12 but 0.3.5 gives this error.

TypeError                                 Traceback (most recent call last)
[<ipython-input-14-16fa2245dec3>](https://localhost:8080/#) in <module>
----> 1 thread_resources.env = ResourceEnv(Mesh(devices, ('a','b')))

TypeError: <lambda>() missing 1 required positional argument: 'loops'

I tried as much as I could imagine but still I can't understand what the 'loops' means.

The whole view is like this

This is the expected result done by JAX0.2.12 below.

I couldn't find any related change in changelog of JAX doc but am I missing something?
Thank you for your help!