jax.interpreters.xla.pytype_aval_mappings is deprecated

[NeilGirdhar]

.venv/lib/python3.13/site-packages/tensorflow_probability/python/internal/backend/jax/ops.py:681
  /home/neil/src/efax/.venv/lib/python3.13/site-packages/tensorflow_probability/python/internal/backend/jax/ops.py:681: DeprecationWarning: jax.interpreters.xla.pytype_aval_mappings is deprecated.

[NeilGirdhar]

@jakevdp Would it be better to skip this code entirely when Jax's version >= 0.5?

[jakevdp]

Yes, you should probably skip. As of 0.5.0, jax.interpreters.xla.pytype_aval_mappings and jax.core.pytype_aval_mappings point to the same object, though the first reference is deprecated. Eventually the first reference will be removed and only jax.core.pytype_aval_mappings will remain.

[NeilGirdhar]

Got it, just wanted to eliminate the warning.