Rewrote hyp2f1 derivatives with scipy functions

Author: ColtAllen

pytensor/scalar/math.py

@@ -1556,15 +1556,15 @@ def _hyp2f1_da(a, b, c, z):
             else:
                 term1 = _infinisum(
                     lambda k: (
-                        (gamma(a + k) / gamma(a))
-                        * (gamma(b + k) / gamma(b))
-                        * psi(a + k)
+                        (scipy.special.gamma(a + k) / scipy.special.gamma(a))
+                        * (scipy.special.gamma(b + k) / scipy.special.gamma(b))
+                        * scipy.special.psi(a + k)
                         * (z**k)
                     )
-                    / (gamma(c + k) / gamma(c))
-                    * gamma(k + 1)
+                    / (scipy.special.gamma(c + k) / scipy.special.gamma(c))
+                    * scipy.special.gamma(k + 1)
                 )
-                term2 = psi(a) * hyp2f1(a, b, c, z)
+                term2 = scipy.special.psi(a) * scipy.special.hyp2f1(a, b, c, z)
 
                 return term1 - term2
 
@@ -1579,15 +1579,15 @@ def _hyp2f1_db(a, b, c, z):
             else:
                 term1 = _infinisum(
                     lambda k: (
-                        (gamma(a + k) / gamma(a))
-                        * (gamma(b + k) / gamma(b))
-                        * psi(b + k)
+                        (scipy.special.gamma(a + k) / scipy.special.gamma(a))
+                        * (scipy.special.gamma(b + k) / scipy.special.gamma(b))
+                        * scipy.special.psi(b + k)
                         * (z**k)
                     )
-                    / (gamma(c + k) / gamma(c))
-                    * gamma(k + 1)
+                    / (scipy.special.gamma(c + k) / scipy.special.gamma(c))
+                    * scipy.special.gamma(k + 1)
                 )
-                term2 = psi(b) * hyp2f1(a, b, c, z)
+                term2 = scipy.special.psi(b) * scipy.special.hyp2f1(a, b, c, z)
 
                 return term1 - term2
 
@@ -1599,16 +1599,16 @@ def _hyp2f1_dc(a, b, c, z):
                 raise NotImplementedError("Gradient not supported for |z| >= 1")
 
             else:
-                term1 = psi(c) * hyp2f1(a, b, c, z)
+                term1 = scipy.special.psi(c) * scipy.special.hyp2f1(a, b, c, z)
                 term2 = _infinisum(
                     lambda k: (
-                        (gamma(a + k) / gamma(a))
-                        * (gamma(b + k) / gamma(b))
-                        * psi(c + k)
+                        (scipy.special.gamma(a + k) / scipy.special.gamma(a))
+                        * (scipy.special.gamma(b + k) / scipy.special.gamma(b))
+                        * scipy.special.psi(c + k)
                         * (z**k)
                     )
-                    / (gamma(c + k) / gamma(c))
-                    * gamma(k + 1)
+                    / (scipy.special.gamma(c + k) / scipy.special.gamma(c))
+                    * scipy.special.gamma(k + 1)
                 )
                 return term1 - term2
 
@@ -1617,7 +1617,7 @@ def _hyp2f1_dz(a, b, c, z):
             Derivative of hyp2f1 wrt z
             """
 
-            return ((a * b) / c) * hyp2f1(a + 1, b + 1, c + 1, z)
+            return ((a * b) / c) * scipy.special.hyp2f1(a + 1, b + 1, c + 1, z)
 
         if wrt == 0:
             return _hyp2f1_da(a, b, c, z)