WIP implement scalar Scan Op

Author: ricardoV94

pytensor/scalar/math.py

@@ -27,13 +27,15 @@
     isinf,
     log,
     log1p,
+    same_out,
     switch,
     true_div,
     upcast,
     upgrade_to_float,
     upgrade_to_float64,
     upgrade_to_float_no_complex,
 )
+from pytensor.scalar.scan import ScalarScanOp
 
 
 class Erf(UnaryScalarOp):
@@ -751,87 +753,180 @@ def c_code(self, *args, **kwargs):
 gammainc_der = GammaIncDer(upgrade_to_float, name="gammainc_der")
 
 
-class GammaIncCDer(BinaryScalarOp):
-    """
-    Gradient of the the regularized upper gamma function (Q) wrt to the first
-    argument (k, a.k.a. alpha). Adapted from STAN `grad_reg_inc_gamma.hpp`
-    """
+# class GammaIncCDer(BinaryScalarOp):
+#     """
+#     Gradient of the the regularized upper gamma function (Q) wrt to the first
+#     argument (k, a.k.a. alpha). Adapted from STAN `grad_reg_inc_gamma.hpp`
+#     """
+#
+#     @staticmethod
+#     def st_impl(k, x):
+#         gamma_k = scipy.special.gamma(k)
+#         digamma_k = scipy.special.digamma(k)
+#         log_x = np.log(x)
+#
+#         # asymptotic expansion http://dlmf.nist.gov/8.11#E2
+#         if (x >= k) and (x >= 8):
+#             S = 0
+#             k_minus_one_minus_n = k - 1
+#             fac = k_minus_one_minus_n
+#             dfac = 1
+#             xpow = x
+#             delta = dfac / xpow
+#
+#             for n in range(1, 10):
+#                 k_minus_one_minus_n -= 1
+#                 S += delta
+#                 xpow *= x
+#                 dfac = k_minus_one_minus_n * dfac + fac
+#                 fac *= k_minus_one_minus_n
+#                 delta = dfac / xpow
+#                 if np.isinf(delta):
+#                     warnings.warn(
+#                         "gammaincc_der did not converge",
+#                         RuntimeWarning,
+#                     )
+#                     return np.nan
+#
+#             return (
+#                 scipy.special.gammaincc(k, x) * (log_x - digamma_k)
+#                 + np.exp(-x + (k - 1) * log_x) * S / gamma_k
+#             )
+#
+#         # gradient of series expansion http://dlmf.nist.gov/8.7#E3
+#         else:
+#             log_precision = np.log(1e-6)
+#             max_iters = int(1e5)
+#             S = 0
+#             log_s = 0.0
+#             s_sign = 1
+#             log_delta = log_s - 2 * np.log(k)
+#             for n in range(1, max_iters + 1):
+#                 S += np.exp(log_delta) if s_sign > 0 else -np.exp(log_delta)
+#                 s_sign = -s_sign
+#                 log_s += log_x - np.log(n)
+#                 log_delta = log_s - 2 * np.log(n + k)
+#
+#                 if np.isinf(log_delta):
+#                     warnings.warn(
+#                         "gammaincc_der did not converge",
+#                         RuntimeWarning,
+#                     )
+#                     return np.nan
+#
+#                 if log_delta <= log_precision:
+#                     return (
+#                         scipy.special.gammainc(k, x) * (digamma_k - log_x)
+#                         + np.exp(k * log_x) * S / gamma_k
+#                     )
+#
+#             warnings.warn(
+#                 f"gammaincc_der did not converge after {n} iterations",
+#                 RuntimeWarning,
+#             )
+#             return np.nan
+#
+#     def impl(self, k, x):
+#         return self.st_impl(k, x)
+#
+#     def c_code(self, *args, **kwargs):
+#         raise NotImplementedError()
+#
+#
+# gammaincc_der = GammaIncCDer(upgrade_to_float, name="gammaincc_der")
+
+
+class GammaIncCDerInnerScan1(ScalarScanOp):
+    nin = 7
+    nout = 6
+    n_steps = 9
+
+    @property
+    def fn(self):
+        def inner_fn(S, delta, xpow, k_minus_one_minus_n, dfac, fac, x):
+            S += delta
+            xpow *= x
+            k_minus_one_minus_n -= 1
+            dfac = k_minus_one_minus_n * dfac + fac
+            fac *= k_minus_one_minus_n
+            delta = dfac / xpow
+            return S, delta, xpow, k_minus_one_minus_n, dfac, fac
 
-    @staticmethod
-    def st_impl(k, x):
-        gamma_k = scipy.special.gamma(k)
-        digamma_k = scipy.special.digamma(k)
-        log_x = np.log(x)
+        return inner_fn
 
-        # asymptotic expansion http://dlmf.nist.gov/8.11#E2
-        if (x >= k) and (x >= 8):
-            S = 0
-            k_minus_one_minus_n = k - 1
-            fac = k_minus_one_minus_n
-            dfac = 1
-            xpow = x
-            delta = dfac / xpow
 
-            for n in range(1, 10):
-                k_minus_one_minus_n -= 1
-                S += delta
-                xpow *= x
-                dfac = k_minus_one_minus_n * dfac + fac
-                fac *= k_minus_one_minus_n
-                delta = dfac / xpow
-                if np.isinf(delta):
-                    warnings.warn(
-                        "gammaincc_der did not converge",
-                        RuntimeWarning,
-                    )
-                    return np.nan
+_gammaincc_der_scan1 = GammaIncCDerInnerScan1(
+    lambda *types: tuple(
+        same_out(type)[0] for type in types[: GammaIncCDerInnerScan1.nout]
+    )
+)
 
+
+class GammaIncCDerInnerScan2(ScalarScanOp):
+    nin = 7
+    nout = 5
+    n_steps = int(1e5)  # maximum number of iterations
+    log_precision = np.log(1e-6)
+
+    @property
+    def fn(self):
+        import pytensor.tensor as pt
+        from pytensor.scan import until
+
+        def inner_fn(S, log_s, s_sign, log_delta, n, k, log_x):
+            delta = pt.exp(log_delta)
+            S += pt.switch(s_sign > 0, delta, -delta)
+            s_sign = -s_sign
+            log_s += log_x - pt.log(n)
+            log_delta = log_s - 2 * pt.log(n + k)
+            n += 1
             return (
-                scipy.special.gammaincc(k, x) * (log_x - digamma_k)
-                + np.exp(-x + (k - 1) * log_x) * S / gamma_k
+                (S, log_s, s_sign, log_delta, n),
+                {},
+                until(pt.all(log_delta < self.log_precision)),
             )
 
-        # gradient of series expansion http://dlmf.nist.gov/8.7#E3
-        else:
-            log_precision = np.log(1e-6)
-            max_iters = int(1e5)
-            S = 0
-            log_s = 0.0
-            s_sign = 1
-            log_delta = log_s - 2 * np.log(k)
-            for n in range(1, max_iters + 1):
-                S += np.exp(log_delta) if s_sign > 0 else -np.exp(log_delta)
-                s_sign = -s_sign
-                log_s += log_x - np.log(n)
-                log_delta = log_s - 2 * np.log(n + k)
-
-                if np.isinf(log_delta):
-                    warnings.warn(
-                        "gammaincc_der did not converge",
-                        RuntimeWarning,
-                    )
-                    return np.nan
-
-                if log_delta <= log_precision:
-                    return (
-                        scipy.special.gammainc(k, x) * (digamma_k - log_x)
-                        + np.exp(k * log_x) * S / gamma_k
-                    )
+        return inner_fn
 
-            warnings.warn(
-                f"gammaincc_der did not converge after {n} iterations",
-                RuntimeWarning,
-            )
-            return np.nan
 
-    def impl(self, k, x):
-        return self.st_impl(k, x)
-
-    def c_code(self, *args, **kwargs):
-        raise NotImplementedError()
+_gammaincc_der_scan2 = GammaIncCDerInnerScan2(
+    lambda *types: tuple(
+        same_out(type)[0] for type in types[: GammaIncCDerInnerScan2.nout]
+    )
+)
 
 
-gammaincc_der = GammaIncCDer(upgrade_to_float, name="gammaincc_der")
+def gammaincc_der(k, x):
+    gamma_k = gamma(k)
+    digamma_k = psi(k)
+    log_x = log(x)
+
+    # asymptotic expansion http://dlmf.nist.gov/8.11#E2
+    S = np.array(0.0, dtype="float64")
+    dfac = np.array(1.0, dtype="float64")
+    xpow = x
+    k_minus_one_minus_n = k - 1
+    fac = k_minus_one_minus_n
+    delta = true_div(dfac, xpow)
+    S, *_ = _gammaincc_der_scan1(S, delta, xpow, k_minus_one_minus_n, fac, dfac, x)
+    res1 = (
+        gammaincc(k, x) * (log_x - digamma_k) + exp(-x + (k - 1) * log_x) * S / gamma_k
+    )
+
+    # gradient of series expansion http://dlmf.nist.gov/8.7#E3
+    S = np.array(0.0, dtype="float64")
+    log_s = np.array(0.0, dtype="float64")
+    s_sign = np.array(1, dtype="int8")
+    n = np.array(1, dtype="int64")
+    log_delta = log_s - 2 * log(k)
+    S, *_ = _gammaincc_der_scan2(S, log_s, s_sign, log_delta, n, k, log_x)
+    res2 = gammainc(k, x) * (digamma_k - log_x) + exp(k * log_x) * S / gamma_k
+
+    return switch(
+        (x >= k) & (x >= 8),
+        res1,
+        res2,
+    )
 
 
 class GammaU(BinaryScalarOp):

pytensor/scalar/scan.py

@@ -0,0 +1,14 @@
+from pytensor.scalar.basic import ScalarOp
+
+
+class ScalarScanOp(ScalarOp):
+    """Dummy Scalar Op that encapsulates a scalar scan operation.
+
+    This Op is never supposed to be evaluated. It can safely be converted
+    to an Elemwise which is rewritten into a Scan node during compilation.
+
+    TODO: FINISH DOCSTRINGS
+    """
+
+    def impl(self, *args, **kwargs):
+        raise RuntimeError("Scalar Scan Ops should never be evaluated!")

pytensor/tensor/rewriting/elemwise.py

@@ -7,7 +7,7 @@
 import pytensor
 import pytensor.scalar.basic as aes
 from pytensor import compile
-from pytensor.compile.mode import get_target_language
+from pytensor.compile.mode import get_target_language, optdb
 from pytensor.configdefaults import config
 from pytensor.graph.basic import Apply, Constant, io_toposort
 from pytensor.graph.features import ReplaceValidate
@@ -20,8 +20,10 @@
 )
 from pytensor.graph.rewriting.db import SequenceDB
 from pytensor.graph.utils import InconsistencyError, MethodNotDefined, TestValueError
+from pytensor.scalar import ScalarScanOp
 from pytensor.tensor.basic import MakeVector, alloc, cast, get_scalar_constant_value
 from pytensor.tensor.elemwise import CAReduce, DimShuffle, Elemwise
+from pytensor.tensor.extra_ops import broadcast_arrays
 from pytensor.tensor.exceptions import NotScalarConstantError
 from pytensor.tensor.rewriting.basic import register_canonicalize, register_specialize
 from pytensor.tensor.shape import shape_padleft
@@ -1025,3 +1027,41 @@ def local_careduce_fusion(fgraph, node):
     "fusion",
     position=49,
 )
+
+
+@node_rewriter([Elemwise])
+def inline_elemwise_scan(fgraph, node):
+    from pytensor.scan.basic import scan
+
+    scalar_op = node.op.scalar_op
+
+    if not isinstance(scalar_op, ScalarScanOp):
+        return None
+
+    # TODO: Add non-batched implementation?
+    n_carried_inputs = scalar_op.nout - scalar_op.nin
+    bcasted_inputs = broadcast_arrays(*node.inputs)
+    ret, updates = scan(
+        scalar_op.fn,
+        outputs_info=bcasted_inputs[:n_carried_inputs],
+        non_sequences=bcasted_inputs[n_carried_inputs:],
+        n_steps=scalar_op.n_steps,
+        sequences=None,
+        strict=True,
+    )
+    if updates:
+        raise ValueError("Scalar scan should never return updates")
+    if scalar_op.nout == 1:
+        ret = (ret,)
+    return [r[-1] for r in ret]
+
+
+# We want to run this before the first merge optimizer
+# and before the first scan optimizer.
+optdb.register(
+    "inline_elemwise_scan",
+    in2out(inline_elemwise_scan),
+    "fast_compile",
+    "fast_run",
+    position=-0.01,
+)

tests/tensor/test_math_scipy.py

@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 
+from pytensor.gradient import verify_grad
 
 scipy = pytest.importorskip("scipy")
 
@@ -9,11 +10,11 @@
 import scipy.special
 import scipy.stats
 
-from pytensor import function
+from pytensor import function, grad
 from pytensor import tensor as at
 from pytensor.compile.mode import get_default_mode
 from pytensor.configdefaults import config
-from pytensor.tensor import inplace
+from pytensor.tensor import inplace, vector, gammaincc
 from tests import unittest_tools as utt
 from tests.tensor.utils import (
     _good_broadcast_unary_chi2sf,
@@ -422,6 +423,23 @@ def test_gammainc_ddk_tabulated_values():
         )
 
 
+def test_gammaincc_ddk_performance(benchmark):
+    rng = np.random.default_rng(1)
+    k = vector("k")
+    x = vector("x")
+
+    out = gammaincc(k, x)
+    grad_fn = function([k, x], grad(out.sum(), wrt=[k]), mode="FAST_RUN")
+    vals = [
+        # Values that hit the second branch of the gradient
+        np.full((1000,), 3.2),
+        np.full((1000,), 0.01),
+    ]
+
+    verify_grad(gammaincc, vals, rng=rng)
+    benchmark(grad_fn, *vals)
+
+
 TestGammaUBroadcast = makeBroadcastTester(
     op=at.gammau,
     expected=expected_gammau,