plot_response_curve doesnt work due to upgrade of jax

[mitsu666]

Hello,

I'm pleased to inform bebefical information. Yesterday (Feb 17th) jax have been upgaraded to version 0.4.4. So jax.vmap function does'nt work well on google colobo.

Therfore I offer resolution here.

After instaling lightweight_mmm, downgrade such a below. !pip uninstall -y jax !pip install jax==0.4.2

!pip uninstall -y jaxlib !pip install jaxlib==0.4.2

I'm sorry that I'm not good at english. Regards,