Moved factorial and poch into tensor.special

Author: ColtAllen

pytensor/scalar/math.py

@@ -33,8 +33,6 @@
     upgrade_to_float,
     upgrade_to_float64,
     upgrade_to_float_no_complex,
-    ScalarType,
-    ScalarVariable,
 )
 
 
@@ -1526,9 +1524,6 @@ class Hyp2F1Der(ScalarOp):
     """
     Derivatives of the Gaussian hypergeometric function ``2F1(a, b; c; z)``.
 
-    Currently written in terms of gamma until poch and factorial Ops are ready:
-    poch(z, m) = (gamma(z + m) / gamma(m))
-    factorial(n) = gamma(n+1)
     """
 
     nin = 5
@@ -1637,20 +1632,4 @@ def c_code(self, *args, **kwargs):
         raise NotImplementedError()
 
 
-hyp2f1_der = Hyp2F1Der(upgrade_to_float, name="hyp2f1_der")
-
-
-def poch(z: ScalarType, m: ScalarType) -> ScalarVariable:
-    """
-    Pochhammer symbol (rising factorial) function.
-
-    """
-    return gamma(z + m) / gamma(z)
-
-
-def factorial(n: ScalarType) -> ScalarVariable:
-    """
-    Factorial function of a scalar or array of numbers.
-
-    """
-    return gamma(n + 1)
+hyp2f1_der = Hyp2F1Der(upgrade_to_float, name="hyp2f1_der")

pytensor/tensor/math.py

@@ -1384,6 +1384,11 @@ def gammal(k, x):
     """Lower incomplete gamma function."""
 
 
+@scalar_elemwise
+def hyp2f1(a, b, c, z):
+    """Gaussian hypergeometric function."""
+
+
 @scalar_elemwise
 def j0(x):
     """Bessel function of the first kind of order 0."""
@@ -1866,21 +1871,6 @@ def clip(x, min, max):
     # for grep: clamp, bound
 
 
-@scalar_elemwise
-def hyp2f1(a, b, c, z):
-    """gaussian hypergeometric function"""
-
-
-@scalar_elemwise
-def poch(z, m):
-    """pochhammer symbol (rising factorial) function"""
-
-
-@scalar_elemwise
-def factorial(n):
-    """factorial function"""
-
-
 pprint.assign(add, printing.OperatorPrinter("+", -2, "either"))
 pprint.assign(mul, printing.OperatorPrinter("*", -1, "either"))
 pprint.assign(sub, printing.OperatorPrinter("-", -2, "left"))
@@ -3148,8 +3138,6 @@ def matmul(x1: "ArrayLike", x2: "ArrayLike", dtype: Optional["DTypeLike"] = None
     "logaddexp",
     "logsumexp",
     "hyp2f1",
-    "poch",
-    "factorial",
 ]
 
 DEPRECATED_NAMES = [

pytensor/tensor/special.py

@@ -1,13 +1,18 @@
 import warnings
 from textwrap import dedent
+from typing import TYPE_CHECKING
 
 import numpy as np
 import scipy
 
 from pytensor.graph.basic import Apply
 from pytensor.link.c.op import COp
 from pytensor.tensor.basic import as_tensor_variable
-from pytensor.tensor.math import neg, sum
+from pytensor.tensor.math import neg, sum, gamma
+
+
+if TYPE_CHECKING:
+    from pytensor.tensor import TensorLike, TensorVariable
 
 
 class SoftmaxGrad(COp):
@@ -768,7 +773,25 @@ def log_softmax(c, axis=UNSET_AXIS):
     return LogSoftmax(axis=axis)(c)
 
 
+def poch(z: "TensorLike", m: "TensorLike") -> "TensorVariable":
+    """
+    Pochhammer symbol (rising factorial) function.
+
+    """
+    return gamma(z + m) / gamma(z)
+
+
+def factorial(n: "TensorLike") -> "TensorVariable":
+    """
+    Factorial function of a scalar or array of numbers.
+
+    """
+    return gamma(n + 1)
+
+
 __all__ = [
     "softmax",
     "log_softmax",
+    "poch",
+    "factorial",
 ]

tests/tensor/test_math_scipy.py

@@ -768,50 +768,18 @@ def test_deprecated_module():
     expected=expected_hyp2f1,
     good=_good_broadcast_quaternary_hyp2f1,
     eps=2e-10,
+    mode=mode_no_scipy,
 )
 
 TestHyp2F1InplaceBroadcast = makeBroadcastTester(
     op=inplace.hyp2f1_inplace,
     expected=expected_hyp2f1,
     good=_good_broadcast_quaternary_hyp2f1,
+    eps=2e-10,
+    mode=mode_no_scipy,
     inplace=True,
 )
 
-_good_broadcast_binary_poch = dict(
-    normal=(
-        random_ranged(0, 5, (2, 1), rng=rng),
-        random_ranged(0, 5, (2, 1), rng=rng),
-    )
-)
-
-
-@pytest.mark.parametrize(
-    "z, m", [random_ranged(0, 5, (2, 1), rng=rng), random_ranged(0, 5, (2, 1), rng=rng)]
-)
-def test_poch(z, m):
-
-    z, m = at.scalars("z", "m")
-
-    poch = at.poch(z, m)
-
-    actual = function([z, m], poch)
-    expected = scipy.special.poch(z, m)
-
-    assert actual == expected
-
-
-@pytest.mark.parametrize("n", random_ranged(0, 5, (2, 1)))
-def test_factorial(n):
-
-    n = at.scalars("n")
-
-    factorial = at.factorial(n)
-
-    actual = function([n], factorial)
-    expected = scipy.special.factorial(n)
-
-    assert actual == expected
-
 
 class TestBetaIncGrad:
     def test_stan_grad_partial(self):

tests/tensor/test_special.py

@@ -1,12 +1,16 @@
 import numpy as np
 import pytest
+from scipy.special import factorial as scipy_factorial
 from scipy.special import log_softmax as scipy_log_softmax
+from scipy.special import poch as scipy_poch
 from scipy.special import softmax as scipy_softmax
 
 from pytensor.compile.function import function
 from pytensor.configdefaults import config
-from pytensor.tensor.special import LogSoftmax, Softmax, SoftmaxGrad, log_softmax, softmax
+from pytensor.tensor import scalar, scalars
+from pytensor.tensor.special import LogSoftmax, Softmax, SoftmaxGrad, log_softmax, softmax, poch, factorial
 from pytensor.tensor.type import matrix, tensor3, tensor4, vector
+from tests.tensor.utils import random_ranged
 from tests import unittest_tools as utt
 
 
@@ -134,3 +138,32 @@ def test_valid_axis(self):
 
         with pytest.raises(ValueError):
             SoftmaxGrad(-4)(*x)
+
+
+@pytest.mark.parametrize(
+    "z, m", [random_ranged(0, 5, (2,)), random_ranged(0, 5, (2,))]
+)
+def test_poch(z, m):
+
+    _z, _m = scalars("z", "m")
+
+    actual_fn = function([_z, _m], poch(_z, _m))
+    actual = actual_fn(z, m)
+
+    expected = scipy_poch(z, m)
+
+    assert np.allclose(actual, expected)
+
+
+@pytest.mark.parametrize("n", random_ranged(0, 5, (1,)))
+def test_factorial(n):
+
+    _n = scalar("n")
+
+    actual_fn = function([_n], factorial(_n))
+    actual = actual_fn(n)
+    
+    expected = scipy_factorial(n)
+
+    assert np.allclose(actual, expected)
+