jax-md renderer not working

[exenGT]

Hi,

I tried running the notebook "NeurIPS Demo.ipynb" in Google Colab. The following section gives an error:

from jax_md.colab_tools import renderer

renderer.render(box_size, 
                {
                    'atom': renderer.Sphere(qm_positions[0]),
                }, 
                resolution=[400, 400])

output error message:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
[<ipython-input-5-92cadf3afdaa>](https://localhost:8080/#) in <module>
----> 1 from jax_md.colab_tools import renderer
      2 
      3 renderer.render(box_size, 
      4                 {
      5                     'atom': renderer.Sphere(qm_positions[0]),

12 frames
[/usr/local/lib/python3.8/dist-packages/jax/_src/core.py](https://localhost:8080/#) in <module>
   2085   return call_bind_continuation, top_trace, fun_, tracers, params
   2086 
-> 2087 @lu.transformation_with_aux
   2088 def process_env_traces_call(primitive: CallPrimitive, level: int,
   2089                             params_tuple: tuple, *args):

AttributeError: partially initialized module 'jax._src.linear_util' has no attribute 'transformation_with_aux' (most likely due to a circular import)

Thanks,
JW

[tzadouri]

Did you fix it?