Add dispatch for Jax Hyp2F1 Function

[ColtAllen]

Description

This PR adds Jax support for the Gaussian Hypergeometric function. This was added just before the latest Jax 0.6.1 release, but tests are still somehow passing for Jax 0.5.2.

Related Issue

• Closes #Add dispatch for numba Hyp2F1 Function #1390
• Related to #

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

---

📚 Documentation preview 📚: https://pytensor--1439.org.readthedocs.build/en/1439/

[ColtAllen]

From the CI run:

AttributeError: module 'jax.scipy.special' has no attribute 'hyp2f1'

Seems we'll need to update the pyproject.toml to minimally support the latest version.

[ricardoV94]

If it passes locally without your PR in a recent enough version of JAX we don't need to do anything, we're not testing most jax.special functions. It's enough if it has the same signature as scipy

[ColtAllen]

If it passes locally without your PR in a recent enough version of JAX we don't need to do anything, we're not testing most jax.special functions. It's enough if it has the same signature as scipy

In pymc-marketing, Nutpie runs endlessly, and Blackjax and NumPyro raise the following error:

~/mamba/envs/pymc-marketing-dev/lib/python3.10/site-packages/jax/_src/lax/control_flow/loops.py:1975, in _while_transpose_error(*_, **kwargs)

ValueError: Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop with dynamic start/stop values. 
Try using lax.scan, or using fori_loop with static start/stop.

Are we at the mercy of the jax developers to resolve this?

[ricardoV94]

Did they use a fori loop in their hyp2f1 implementation? If so you may want to open an issue with them to use scan instead. But if they used a while loop you're out of luck, that can't be differentiated no matter what.

You may want to try nutpie with the numba backend after #1445

[ColtAllen]

Did they use a fori loop in their hyp2f1 implementation? If so you may want to open an issue with them to use scan instead. But if they used a while loop you're out of luck, that can't be differentiated no matter what.

lax.while_loop calls are returned by _hyp2f1_serie, a conditional called when all parameters are > 0. The parameters in question for ParetoNBDModel are Gamma shape & scale params, so we're out of luck 🙁

You may want to try nutpie with the numba backend after #1445

numba supports scipy.special.hyp2f1; hope is kindled!

[ColtAllen]

Referenced issue has been updated for numba.