fix module 'jax' has no attribute 'tree_multimap'

[Quandong-Zhang]

Set the jax package to the correct version.
Fixed issue with #71

[Quandong-Zhang]

I found that the first change does not seem to solve the problem,so i fix it again.

#Install the correct version of the jax package
!pip install jax==0.3.15

The placement of these lines of code doesn't seem to be correct in the first change.
It was placed in a relatively advanced position, and I guess it was caused by the installation of the dependent library later.

[a1764983904]

I can't use, did not appear each time you run into similar to the module