Add Exponentially scaled modified Bessel Op

[dehorsley]

Fixes #542

As with Iv, uses scipy implementation in C code and tensorflow probability implementation for jax.

Type of change

• New feature / enhancement

[codecov-commenter]

Codecov Report

Merging #543 (b032e25) into main (68b41a4) will increase coverage by 0.00%.
The diff coverage is 81.81%.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #543   +/-   ##
=======================================
  Coverage   80.90%   80.90%           
=======================================
  Files         162      162           
  Lines       46393    46415   +22     
  Branches    11349    11353    +4     
=======================================
+ Hits        37535    37553   +18     
- Misses       6635     6639    +4     
  Partials     2223     2223           

Files	Coverage Δ
pytensor/tensor/inplace.py	100.00% <100.00%> (ø)
pytensor/tensor/math.py	90.08% <100.00%> (+0.01%)	⬆️
pytensor/link/jax/dispatch/scalar.py	93.63% <50.00%> (-1.15%)	⬇️
pytensor/scalar/math.py	88.92% <85.71%> (-0.06%)	⬇️

[ricardoV94]

Looks great, could also be a nice touch to add a rewrite to convert the iv/exp form into ive so users writing the old form get the more stable one now automatically

[ricardoV94]

Could use PatternRewriter like we did here: #505

[dehorsley]

Looks great, could also be a nice touch to add a rewrite to convert the iv/exp form into ive so users writing the old form get the more stable one now automatically

Sounds good. The tricky thing is that we need an abs in there unless we can prove the argument is positive and real. By which I mean the rewrite is:

iv(nu, x)/exp(abs(x)) -> ive(nu, x)

but the use in PyMC is:

a = 1 / pt.square(self.ls)
c = pt.where(J > 0, 2, 1)

q2 = c * pt.iv(J, a) / pt.exp(a)

So unless we carry some "positive real" tag(s) around, we won't pickup this kind of case. Maybe not worth the complexity, and just add the above rewrite?

Also, do you have a suggestion on how to handle the two forms? I.e. can I catch both these as one rewrite:

iv(nu, x)/exp(abs(x)) -> ive(nu, x)
iv(nu, x)*exp(-abs(x)) -> ive(nu, x)

Is there a canonicalizing rewrite that will handle that?

[ricardoV94]

Canonicalize wouldn't do that kind of modification. You can write a more typical rewrite with a function that handles both cases.

But since it needs to be positive maybe not worth adding for now.

The kind of machinery in #525 would make it really easy to complement rewrites with positivity/negative analysis. This case may be a nice motivating example.

[dehorsley]

But since it needs to be positive maybe not worth adding for now.

So happy to go as is, or would you like me to add the rewrite with the abs in there?

[ricardoV94]

Happy as is