Jax: scipy version pin

[gante]

What does this PR do?

Pins the scipy version so as to avoid jax import errors. Our pinned jax versions import (e.g.) scipy.linalg.tril, which was deprecated -- search for tril in this release page.

Example of failing import (from a fresh venv, which was getting scipy==1.13.0):

tests/models/t5/test_modeling_flax_t5.py:45: in <module>
    import optax
../venvs/hf/lib/python3.10/site-packages/optax/__init__.py:18: in <module>
    from optax._src.alias import adabelief
../venvs/hf/lib/python3.10/site-packages/optax/_src/alias.py:24: in <module>
    from optax._src import factorized
../venvs/hf/lib/python3.10/site-packages/optax/_src/factorized.py:27: in <module>
    from optax._src import utils
../venvs/hf/lib/python3.10/site-packages/optax/_src/utils.py:22: in <module>
    import jax.scipy.stats.norm as multivariate_normal                                                                                                ../venvs/hf/lib/python3.10/site-packages/jax/scipy/stats/__init__.py:39: in <module>
    from jax._src.scipy.stats.kde import gaussian_kde as gaussian_kde
../venvs/hf/lib/python3.10/site-packages/jax/_src/scipy/stats/kde.py:26: in <module>
    from jax.scipy import linalg, special
../venvs/hf/lib/python3.10/site-packages/jax/scipy/linalg.py:18: in <module>
    from jax._src.scipy.linalg import (
../venvs/hf/lib/python3.10/site-packages/jax/_src/scipy/linalg.py:408: in <module>
    @_wraps(scipy.linalg.tril)
E   AttributeError: module 'scipy.linalg' has no attribute 'tril'

[amyeroberts]

Thanks for handling!

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.