More direct access to special functions

Author: ricardoV94

pytensor/scalar/math.py

@@ -9,7 +9,7 @@
 from textwrap import dedent
 
 import numpy as np
-import scipy.special
+from scipy import special
 
 from pytensor.configdefaults import config
 from pytensor.gradient import grad_not_implemented, grad_undefined
@@ -53,7 +53,7 @@ class Erf(UnaryScalarOp):
     nfunc_spec = ("scipy.special.erf", 1, 1)
 
     def impl(self, x):
-        return scipy.special.erf(x)
+        return special.erf(x)
 
     def L_op(self, inputs, outputs, grads):
         (x,) = inputs
@@ -87,7 +87,7 @@ class Erfc(UnaryScalarOp):
     nfunc_spec = ("scipy.special.erfc", 1, 1)
 
     def impl(self, x):
-        return scipy.special.erfc(x)
+        return special.erfc(x)
 
     def L_op(self, inputs, outputs, grads):
         (x,) = inputs
@@ -114,7 +114,7 @@ def c_code(self, node, name, inp, out, sub):
         return f"{z} = erfc(({cast}){x});"
 
 
-# scipy.special.erfc don't support complex. Why?
+# special.erfc don't support complex. Why?
 erfc = Erfc(upgrade_to_float_no_complex, name="erfc")
 
 
@@ -136,7 +136,7 @@ class Erfcx(UnaryScalarOp):
     nfunc_spec = ("scipy.special.erfcx", 1, 1)
 
     def impl(self, x):
-        return scipy.special.erfcx(x)
+        return special.erfcx(x)
 
     def L_op(self, inputs, outputs, grads):
         (x,) = inputs
@@ -192,7 +192,7 @@ class Erfinv(UnaryScalarOp):
     nfunc_spec = ("scipy.special.erfinv", 1, 1)
 
     def impl(self, x):
-        return scipy.special.erfinv(x)
+        return special.erfinv(x)
 
     def L_op(self, inputs, outputs, grads):
         (x,) = inputs
@@ -227,7 +227,7 @@ class Erfcinv(UnaryScalarOp):
     nfunc_spec = ("scipy.special.erfcinv", 1, 1)
 
     def impl(self, x):
-        return scipy.special.erfcinv(x)
+        return special.erfcinv(x)
 
     def L_op(self, inputs, outputs, grads):
         (x,) = inputs
@@ -262,7 +262,7 @@ class Owens_t(BinaryScalarOp):
     nfunc_spec = ("scipy.special.owens_t", 2, 1)
 
     def impl(self, h, a):
-        return scipy.special.owens_t(h, a)
+        return special.owens_t(h, a)
 
     def grad(self, inputs, grads):
         (h, a) = inputs
@@ -287,7 +287,7 @@ class Gamma(UnaryScalarOp):
     nfunc_spec = ("scipy.special.gamma", 1, 1)
 
     def impl(self, x):
-        return scipy.special.gamma(x)
+        return special.gamma(x)
 
     def L_op(self, inputs, outputs, gout):
         (x,) = inputs
@@ -322,7 +322,7 @@ class GammaLn(UnaryScalarOp):
     nfunc_spec = ("scipy.special.gammaln", 1, 1)
 
     def impl(self, x):
-        return scipy.special.gammaln(x)
+        return special.gammaln(x)
 
     def L_op(self, inputs, outputs, grads):
         (x,) = inputs
@@ -362,7 +362,7 @@ class Psi(UnaryScalarOp):
     nfunc_spec = ("scipy.special.psi", 1, 1)
 
     def impl(self, x):
-        return scipy.special.psi(x)
+        return special.psi(x)
 
     def L_op(self, inputs, outputs, grads):
         (x,) = inputs
@@ -449,7 +449,7 @@ class TriGamma(UnaryScalarOp):
     """
 
     def impl(self, x):
-        return scipy.special.polygamma(1, x)
+        return special.polygamma(1, x)
 
     def L_op(self, inputs, outputs, outputs_gradients):
         (x,) = inputs
@@ -548,7 +548,7 @@ def output_types_preference(n_type, x_type):
         return upgrade_to_float_no_complex(x_type)
 
     def impl(self, n, x):
-        return scipy.special.polygamma(n, x)
+        return special.polygamma(n, x)
 
     def L_op(self, inputs, outputs, output_gradients):
         (n, x) = inputs
@@ -573,10 +573,10 @@ class Chi2SF(BinaryScalarOp):
         ie. chi2 pvalue (chi2 'survival function')
     """
 
-    nfunc_spec = ("scipy.stats.chi2.sf", 2, 1)
+    nfunc_spec = ("scipy.special.chdtrc", 2, 1)
 
     def impl(self, x, k):
-        return scipy.stats.chi2.sf(x, k)
+        return special.chdtrc(x, k)
 
     def c_support_code(self, **kwargs):
         return (C_CODE_PATH / "gamma.c").read_text(encoding="utf-8")
@@ -615,7 +615,7 @@ class GammaInc(BinaryScalarOp):
     nfunc_spec = ("scipy.special.gammainc", 2, 1)
 
     def impl(self, k, x):
-        return scipy.special.gammainc(k, x)
+        return special.gammainc(k, x)
 
     def grad(self, inputs, grads):
         (k, x) = inputs
@@ -662,7 +662,7 @@ class GammaIncC(BinaryScalarOp):
     nfunc_spec = ("scipy.special.gammaincc", 2, 1)
 
     def impl(self, k, x):
-        return scipy.special.gammaincc(k, x)
+        return special.gammaincc(k, x)
 
     def grad(self, inputs, grads):
         (k, x) = inputs
@@ -709,7 +709,7 @@ class GammaIncInv(BinaryScalarOp):
     nfunc_spec = ("scipy.special.gammaincinv", 2, 1)
 
     def impl(self, k, x):
-        return scipy.special.gammaincinv(k, x)
+        return special.gammaincinv(k, x)
 
     def grad(self, inputs, grads):
         (k, x) = inputs
@@ -734,7 +734,7 @@ class GammaIncCInv(BinaryScalarOp):
     nfunc_spec = ("scipy.special.gammainccinv", 2, 1)
 
     def impl(self, k, x):
-        return scipy.special.gammainccinv(k, x)
+        return special.gammainccinv(k, x)
 
     def grad(self, inputs, grads):
         (k, x) = inputs
@@ -969,7 +969,7 @@ class GammaU(BinaryScalarOp):
     # Note there is no basic SciPy version so no nfunc_spec.
 
     def impl(self, k, x):
-        return scipy.special.gammaincc(k, x) * scipy.special.gamma(k)
+        return special.gammaincc(k, x) * special.gamma(k)
 
     def c_support_code(self, **kwargs):
         return (C_CODE_PATH / "gamma.c").read_text(encoding="utf-8")
@@ -1001,7 +1001,7 @@ class GammaL(BinaryScalarOp):
     # Note there is no basic SciPy version so no nfunc_spec.
 
     def impl(self, k, x):
-        return scipy.special.gammainc(k, x) * scipy.special.gamma(k)
+        return special.gammainc(k, x) * special.gamma(k)
 
     def c_support_code(self, **kwargs):
         return (C_CODE_PATH / "gamma.c").read_text(encoding="utf-8")
@@ -1033,7 +1033,7 @@ class Jv(BinaryScalarOp):
     nfunc_spec = ("scipy.special.jv", 2, 1)
 
     def impl(self, v, x):
-        return scipy.special.jv(v, x)
+        return special.jv(v, x)
 
     def grad(self, inputs, grads):
         v, x = inputs
@@ -1058,7 +1058,7 @@ class J1(UnaryScalarOp):
     nfunc_spec = ("scipy.special.j1", 1, 1)
 
     def impl(self, x):
-        return scipy.special.j1(x)
+        return special.j1(x)
 
     def grad(self, inputs, grads):
         (x,) = inputs
@@ -1085,7 +1085,7 @@ class J0(UnaryScalarOp):
     nfunc_spec = ("scipy.special.j0", 1, 1)
 
     def impl(self, x):
-        return scipy.special.j0(x)
+        return special.j0(x)
 
     def grad(self, inp, grads):
         (x,) = inp
@@ -1112,7 +1112,7 @@ class Iv(BinaryScalarOp):
     nfunc_spec = ("scipy.special.iv", 2, 1)
 
     def impl(self, v, x):
-        return scipy.special.iv(v, x)
+        return special.iv(v, x)
 
     def grad(self, inputs, grads):
         v, x = inputs
@@ -1137,7 +1137,7 @@ class I1(UnaryScalarOp):
     nfunc_spec = ("scipy.special.i1", 1, 1)
 
     def impl(self, x):
-        return scipy.special.i1(x)
+        return special.i1(x)
 
     def grad(self, inputs, grads):
         (x,) = inputs
@@ -1159,7 +1159,7 @@ class I0(UnaryScalarOp):
     nfunc_spec = ("scipy.special.i0", 1, 1)
 
     def impl(self, x):
-        return scipy.special.i0(x)
+        return special.i0(x)
 
     def grad(self, inp, grads):
         (x,) = inp
@@ -1181,7 +1181,7 @@ class Ive(BinaryScalarOp):
     nfunc_spec = ("scipy.special.ive", 2, 1)
 
     def impl(self, v, x):
-        return scipy.special.ive(v, x)
+        return special.ive(v, x)
 
     def grad(self, inputs, grads):
         v, x = inputs
@@ -1206,7 +1206,7 @@ class Kve(BinaryScalarOp):
     nfunc_spec = ("scipy.special.kve", 2, 1)
 
     def impl(self, v, x):
-        return scipy.special.kve(v, x)
+        return special.kve(v, x)
 
     def L_op(self, inputs, outputs, output_grads):
         v, x = inputs
@@ -1236,7 +1236,7 @@ class Sigmoid(UnaryScalarOp):
     nfunc_spec = ("scipy.special.expit", 1, 1)
 
     def impl(self, x):
-        return scipy.special.expit(x)
+        return special.expit(x)
 
     def grad(self, inp, grads):
         (x,) = inp
@@ -1403,7 +1403,7 @@ class BetaInc(ScalarOp):
     nfunc_spec = ("scipy.special.betainc", 3, 1)
 
     def impl(self, a, b, x):
-        return scipy.special.betainc(a, b, x)
+        return special.betainc(a, b, x)
 
     def grad(self, inp, grads):
         a, b, x = inp
@@ -1663,7 +1663,7 @@ class BetaIncInv(ScalarOp):
     nfunc_spec = ("scipy.special.betaincinv", 3, 1)
 
     def impl(self, a, b, x):
-        return scipy.special.betaincinv(a, b, x)
+        return special.betaincinv(a, b, x)
 
     def grad(self, inputs, grads):
         (a, b, x) = inputs
@@ -1702,7 +1702,7 @@ class Hyp2F1(ScalarOp):
     nfunc_spec = ("scipy.special.hyp2f1", 4, 1)
 
     def impl(self, a, b, c, z):
-        return scipy.special.hyp2f1(a, b, c, z)
+        return special.hyp2f1(a, b, c, z)
 
     def grad(self, inputs, grads):
         a, b, c, z = inputs