JAXLIB not found when importing optax

[Boyi-Zhang]

Hi,

I am a beginner to Tidy3D.
I am trying to repeat the inverse design from the example notebooks.

however, when importing optax it reports the missing of jaxlib module. this seems due to the switch from jax to autodiff.
As suggested by the website, i tried pip tidy3d[jax] in the website notebook terminal but not working.

What should I do?
THX in advance.

Best regards,
BOyi

[Boyi-Zhang]

I got answer from flexcompute:
!pip install jax
!pip install jaxlib==0.4.26