Provide JAX Ops from Optional tensorflow-probability dependency

[ricardoV94]

Closes #43
Closes #91
Closes #256

[codecov-commenter]

Codecov Report

Merging #403 (e4766a8) into main (8ac8342) will increase coverage by 0.00%.
The diff coverage is 90.90%.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #403   +/-   ##
=======================================
  Coverage   80.46%   80.47%           
=======================================
  Files         156      156           
  Lines       45558    45579   +21     
  Branches    11162    11167    +5     
=======================================
+ Hits        36660    36679   +19     
- Misses       6695     6697    +2     
  Partials     2203     2203           

Files Changed	Coverage Δ
pytensor/link/jax/dispatch/scalar.py	97.27% <90.90%> (-1.14%)	⬇️

[jessegrabowski]

LGTM, seems very straight-forward. Does anything else need to be done after this is merged to JAX compile models with distributions that rely on these functions?

[ricardoV94]

Nope, users just have to install tfp, but they will get the informative message if that's the case

[maresb]

Do you need a version pin in the PyTensor feedstock? You can add a run_constrained which says "if TFP is installed, then make sure it satisfies this version spec".

[ricardoV94]

I don't think so, we didn't need it for numpyro? We can wait and see if people have issues with it?