tfp and jax version issue

[b-biswas]

Jax-Galsim import fails with jax 0.7 due to package compatibility issues.
Here is the stack trace of the error:

    import jax_galsim
../../../micromamba/envs/dfmdet/lib/python3.11/site-packages/jax_galsim/__init__.py:55: in <module>
    from .moffat import Moffat
../../../micromamba/envs/dfmdet/lib/python3.11/site-packages/jax_galsim/moffat.py:7: in <module>
    from jax_galsim.bessel import j0, kv
../../../micromamba/envs/dfmdet/lib/python3.11/site-packages/jax_galsim/bessel.py:4: in <module>
    from tensorflow_probability.substrates.jax.math import bessel_kve as _tfp_bessel_kve
../../../micromamba/envs/dfmdet/lib/python3.11/site-packages/tensorflow_probability/substrates/jax/__init__.py:42: in <module>
    from tensorflow_probability.substrates.jax import bijectors
../../../micromamba/envs/dfmdet/lib/python3.11/site-packages/tensorflow_probability/substrates/jax/bijectors/__init__.py:19: in <module>
    from tensorflow_probability.substrates.jax.bijectors.absolute_value import AbsoluteValue
../../../micromamba/envs/dfmdet/lib/python3.11/site-packages/tensorflow_probability/substrates/jax/bijectors/absolute_value.py:17: in <module>
    from tensorflow_probability.python.internal.backend.jax.compat import v2 as tf
../../../micromamba/envs/dfmdet/lib/python3.11/site-packages/tensorflow_probability/python/internal/backend/jax/__init__.py:19: in <module>
    from tensorflow_probability.python.internal.backend.jax import compat
../../../micromamba/envs/dfmdet/lib/python3.11/site-packages/tensorflow_probability/python/internal/backend/jax/compat.py:17: in <module>
    from tensorflow_probability.python.internal.backend.jax import v1
../../../micromamba/envs/dfmdet/lib/python3.11/site-packages/tensorflow_probability/python/internal/backend/jax/v1.py:23: in <module>
    from tensorflow_probability.python.internal.backend.jax import linalg_impl
../../../micromamba/envs/dfmdet/lib/python3.11/site-packages/tensorflow_probability/python/internal/backend/jax/linalg_impl.py:23: in <module>
    from tensorflow_probability.python.internal.backend.jax import ops
../../../micromamba/envs/dfmdet/lib/python3.11/site-packages/tensorflow_probability/python/internal/backend/jax/ops.py:681: in <module>
    jax.interpreters.xla.pytype_aval_mappings[onp.ndarray])
../../../micromamba/envs/dfmdet/lib/python3.11/site-packages/jax/_src/deprecations.py:54: in getattr
    raise AttributeError(message)
E   AttributeError: jax.interpreters.xla.pytype_aval_mappings was deprecated in JAX v0.5.0 and removed in JAX v0.7.0. jax.core.pytype_aval_mappings can be used as a replacement in most cases.```

[ismael-mendoza]

seems to be tensorflow using deprecated jax features: tensorflow/probability#2010

[ismael-mendoza]

apparently we should use tfp-nightly tensorflow/probability#1994. Specifically: tensorflow/probability#1994 (comment)

[beckermr]

Having people pin to a nightly build is kind of insane. I'd rather pin jax