Fix errors in notebook due to version update

[HarshvirSandhu]

Update PyMC functions

Update wrapping_jax_function notebook for compatibility with newer versions of PyMC

• Fixes: wrapping jax example code cell 22 returns error #625

---

📚 Documentation preview 📚: https://pymc-examples--630.org.readthedocs.build/en/630/

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[twiecki]

This is great, but you need to run pre-commit.

[HarshvirSandhu]

@twiecki The pre-commits are passing locally, but there is an error during build after I push changes

black-jupyter.................................................................................Passed
nbqa-isort....................................................................................Passed
nbqa-pyupgrade................................................................................Passed
Check cells were executed sequentially........................................................Passed
bibtex-tidy...............................................................(no files to check)Skipped
Check notebooks have watermark (see Jupyter style guide from PyMC docs).......................Passed
Check notebooks do not use literally notebook_name as target..................................Passed
Check no references that should be sphinx cross-references are urls...........................Passed
jupytext......................................................................................Passed