Couldn't Run Tutorial Examples on Google Colab

[Xiphofish]

Hello,
I was trying to run the tutorial examples on google colab but I got an error as below when I import the pybamm...
After restart runtime, I still got the same error. Did I miss anything?
Thank you for your help!

ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
albumentations 0.1.12 requires imgaug<0.2.7,>=0.2.5, but you have imgaug 0.2.9 which is incompatible.

AttributeError Traceback (most recent call last)
in ()
3
4 get_ipython().magic('pip install pybamm -q # install PyBaMM if it is not installed')
----> 5 import pybamm

1 frames
/usr/local/lib/python3.7/dist-packages/pybamm/solvers/jax_bdf_solver.py in ()
66 return caller
67
---> 68 @jax.partial(jax.custom_vjp, nondiff_argnums=(0, 1, 2, 3))
69 def _bdf_odeint(fun, mass, rtol, atol, y0, t_eval, *args):
70 """

AttributeError: module 'jax' has no attribute 'partial'

[Saransh-cpp]

@Xiphofish this should be fixed in the next release (PR #1874). For now, you can replace %pip install pybamm -q with -

%pip uninstall --yes jax
%pip install pybamm --no-cache -q

[Xiphofish]

Dear @Saransh-cpp
Thanks for the solution. It looks like it works.

After uninstall the jax, there is still an error though - ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts. albumentations 0.1.12 requires imgaug<0.2.7,>=0.2.5, but you have imgaug 0.2.9 which is incompatible..

Looking forward to the next release!
Thanks!

[Saransh-cpp]

Great! You can ignore that error. The cause again is the Google Colab environment, but it should not affect the usage of PyBaMM.