pytensor 2.24.1__tar.gz → 2.24.2__tar.gz

{pytensor-2.24.1/pytensor.egg-info → pytensor-2.24.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pytensor
-Version: 2.24.1
+Version: 2.24.2
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License: .. _license:
@@ -94,6 +94,7 @@ Requires-Dist: pytest-cov>=2.6.1; extra == "tests"
 Requires-Dist: coverage>=5.1; extra == "tests"
 Requires-Dist: pytest-benchmark; extra == "tests"
 Requires-Dist: pytest-mock; extra == "tests"
+Requires-Dist: pytest-sphinx; extra == "tests"
 Provides-Extra: rtd
 Requires-Dist: sphinx<6,>=5.1.0; extra == "rtd"
 Requires-Dist: pygments; extra == "rtd"

{pytensor-2.24.1 → pytensor-2.24.2}/doc/tutorial/gradients.rst

@@ -267,6 +267,16 @@ or, making use of the R-operator:
 >>> f([4, 4], [2, 2])
 array([ 4.,  4.])
 
+There is a builtin helper that uses the first method
+
+>>> x = pt.dvector('x')
+>>> v = pt.dvector('v')
+>>> y = pt.sum(x ** 2)
+>>> Hv = pytensor.gradient.hessian_vector_product(y, x, v)
+>>> f = pytensor.function([x, v], Hv)
+>>> f([4, 4], [2, 2])
+array([ 4.,  4.])
+
 
 Final Pointers
 ==============

{pytensor-2.24.1 → pytensor-2.24.2}/pyproject.toml

@@ -75,6 +75,7 @@ tests = [
     "coverage>=5.1",
     "pytest-benchmark",
     "pytest-mock",
+    "pytest-sphinx",
 ]
 rtd = ["sphinx>=5.1.0,<6", "pygments", "pydot", "pydot2", "pydot-ng"]
 jax = ["jax", "jaxlib"]
@@ -116,7 +117,7 @@ versionfile_build = "pytensor/_version.py"
 tag_prefix = "rel-"
 
 [tool.pytest]
-addopts = "--durations=50"
+addopts = "--durations=50  --doctest-modules pytensor --ignore=pytensor/misc/check_duplicate_key.py --ignore=pytensor/link"
 testpaths = "tests/"
 
 [tool.ruff]

{pytensor-2.24.1 → pytensor-2.24.2}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-07-04T12:10:31+0200",
+ "date": "2024-07-08T14:11:01+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "afc1a6caf595cb1bf08c9cb1901e1b9e0bb98efc",
- "version": "2.24.1"
+ "full-revisionid": "7f623fefbed89442737fd6d8ba11c32d2eab65ed",
+ "version": "2.24.2"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.24.1 → pytensor-2.24.2}/pytensor/gradient.py

@@ -2050,6 +2050,85 @@ def hessian(cost, wrt, consider_constant=None, disconnected_inputs="raise"):
     return as_list_or_tuple(using_list, using_tuple, hessians)
 
 
+def hessian_vector_product(cost, wrt, p, **grad_kwargs):
+    """Return the expression of the Hessian times a vector p.
+
+    Notes
+    -----
+    This function uses backward autodiff twice to obtain the desired expression.
+    You may want to manually build the equivalent expression by combining backward
+    followed by forward (if all Ops support it) autodiff.
+    See {ref}`docs/_tutcomputinggrads#Hessian-times-a-Vector` for how to do this.
+
+    Parameters
+    ----------
+    cost: Scalar (0-dimensional) variable.
+    wrt: Vector (1-dimensional tensor) 'Variable' or list of Vectors
+    p: Vector (1-dimensional tensor) 'Variable' or list of Vectors
+        Each vector will be used for the hessp wirt to exach input variable
+    **grad_kwargs:
+        Keyword arguments passed to `grad` function.
+
+    Returns
+    -------
+    :class:` Vector or list of Vectors
+        The Hessian times p of the `cost` with respect to (elements of) `wrt`.
+
+    Examples
+    --------
+
+    .. testcode::
+
+        import numpy as np
+        from scipy.optimize import minimize
+        from pytensor import function
+        from pytensor.tensor import vector
+        from pytensor.gradient import grad, hessian_vector_product
+
+        x = vector('x')
+        p = vector('p')
+
+        rosen = (100 * (x[1:] - x[:-1] ** 2) ** 2 + (1 - x[:-1]) ** 2).sum()
+        rosen_jac = grad(rosen, x)
+        rosen_hessp = hessian_vector_product(rosen, x, p)
+
+        rosen_fn = function([x], rosen)
+        rosen_jac_fn = function([x], rosen_jac)
+        rosen_hessp_fn = function([x, p], rosen_hessp)
+        x0 = np.array([1.3, 0.7, 0.8, 1.9, 1.2])
+        res = minimize(
+            rosen_fn,
+            x0,
+            method="Newton-CG",
+            jac=rosen_jac_fn,
+            hessp=rosen_hessp_fn,
+            options={"xtol": 1e-8},
+        )
+        print(res.x)
+
+    .. testoutput::
+
+        [1.         1.         1.         0.99999999 0.99999999]
+
+
+
+    """
+    wrt_list = wrt if isinstance(wrt, Sequence) else [wrt]
+    p_list = p if isinstance(p, Sequence) else [p]
+    grad_wrt_list = grad(cost, wrt=wrt_list, **grad_kwargs)
+    hessian_cost = pytensor.tensor.add(
+        *[
+            (grad_wrt * p).sum()
+            for grad_wrt, p in zip(grad_wrt_list, p_list, strict=True)
+        ]
+    )
+    Hp_list = grad(hessian_cost, wrt=wrt_list, **grad_kwargs)
+
+    if isinstance(wrt, Variable):
+        return Hp_list[0]
+    return Hp_list
+
+
 def _is_zero(x):
     """
     Returns 'yes', 'no', or 'maybe' indicating whether x
@@ -2241,7 +2320,7 @@ def grad_clip(x, lower_bound, upper_bound):
     >>> z2 = pytensor.gradient.grad(x**2, x)
     >>> f = pytensor.function([x], outputs = [z, z2])
     >>> print(f(2.0))
-    [array(1.0), array(4.0)]
+    [array(1.), array(4.)]
 
     Notes
     -----

{pytensor-2.24.1 → pytensor-2.24.2}/pytensor/graph/basic.py

@@ -1034,7 +1034,10 @@ def orphans_between(
 
     Examples
     --------
-    >>> orphans_between([x], [(x+y).out])
+    >>> from pytensor.graph.basic import orphans_between
+    >>> from pytensor.tensor import scalars
+    >>> x, y = scalars("xy")
+    >>> list(orphans_between([x], [(x+y)]))
     [y]
 
     """

{pytensor-2.24.1 → pytensor-2.24.2}/pytensor/link/jax/dispatch/nlinalg.py

@@ -1,4 +1,5 @@
 import jax.numpy as jnp
+import numpy as np
 
 from pytensor.link.jax.dispatch import jax_funcify
 from pytensor.tensor.blas import BatchedDot
@@ -137,12 +138,10 @@ def jax_funcify_Argmax(op, **kwargs):
 
         # NumPy does not support multiple axes for argmax; this is a
         # work-around
-        keep_axes = jnp.array(
-            [i for i in range(x.ndim) if i not in axes], dtype="int64"
-        )
+        keep_axes = np.array([i for i in range(x.ndim) if i not in axes], dtype="int64")
         # Not-reduced axes in front
         transposed_x = jnp.transpose(
-            x, jnp.concatenate((keep_axes, jnp.array(axes, dtype="int64")))
+            x, tuple(np.concatenate((keep_axes, np.array(axes, dtype="int64"))))
         )
         kept_shape = transposed_x.shape[: len(keep_axes)]
         reduced_shape = transposed_x.shape[len(keep_axes) :]
@@ -151,9 +150,9 @@ def jax_funcify_Argmax(op, **kwargs):
         # Otherwise reshape would complain citing float arg
         new_shape = (
             *kept_shape,
-            jnp.prod(jnp.array(reduced_shape, dtype="int64"), dtype="int64"),
+            np.prod(np.array(reduced_shape, dtype="int64"), dtype="int64"),
         )
-        reshaped_x = transposed_x.reshape(new_shape)
+        reshaped_x = transposed_x.reshape(tuple(new_shape))
 
         max_idx_res = jnp.argmax(reshaped_x, axis=-1).astype("int64")
 

{pytensor-2.24.1 → pytensor-2.24.2}/pytensor/link/pytorch/dispatch/basic.py

@@ -6,7 +6,7 @@ from pytensor.compile.ops import DeepCopyOp
 from pytensor.graph.fg import FunctionGraph
 from pytensor.link.utils import fgraph_to_python
 from pytensor.raise_op import CheckAndRaise
-from pytensor.tensor.basic import Alloc, AllocEmpty, ARange, Join
+from pytensor.tensor.basic import Alloc, AllocEmpty, ARange, Eye, Join
 
 
 @singledispatch
@@ -100,3 +100,19 @@ def pytorch_funcify_Join(op, **kwargs):
         return torch.cat(tensors, dim=axis)
 
     return join
+
+
+@pytorch_funcify.register(Eye)
+def pytorch_funcify_eye(op, **kwargs):
+    torch_dtype = getattr(torch, op.dtype)
+
+    def eye(N, M, k):
+        major, minor = (M, N) if k > 0 else (N, M)
+        k_abs = torch.abs(k)
+        zeros = torch.zeros(N, M, dtype=torch_dtype)
+        if k_abs < major:
+            l_ones = torch.min(major - k_abs, minor)
+            return zeros.diagonal_scatter(torch.ones(l_ones, dtype=torch_dtype), k)
+        return zeros
+
+    return eye

{pytensor-2.24.1 → pytensor-2.24.2}/pytensor/misc/pkl_utils.py

@@ -239,7 +239,7 @@ def dump(
     >>> foo_2 = pytensor.shared(1, name='foo')
     >>> with open('model.zip', 'wb') as f:
     ...     dump((foo_1, foo_2, np.array(2)), f)
-    >>> np.load('model.zip').keys()
+    >>> list(np.load('model.zip').keys())
     ['foo', 'foo_2', 'array_0', 'pkl']
     >>> np.load('model.zip')['foo']
     array(0)

{pytensor-2.24.1 → pytensor-2.24.2}/pytensor/scalar/basic.py

@@ -208,6 +208,7 @@ class autocast_float_as:
 
     Examples
     --------
+    >>> from pytensor.tensor import fvector
     >>> with autocast_float_as('float32'):
     ...    assert (fvector() + 1.1).dtype == 'float32'  # temporary downcasting
     >>> assert (fvector() + 1.1).dtype == 'float64' # back to default behaviour

pytensor-2.24.2/pytensor/scalar/c_code/incbet.c

@@ -0,0 +1,311 @@
+/*    adapted from file incbet.c, obtained from the Cephes library (MIT License)
+Cephes Math Library, Release 2.8:  June, 2000
+Copyright 1984, 1995, 2000 by Stephen L. Moshier
+*/
+
+//For GPU support
+#ifdef __CUDACC__
+#define DEVICE __device__
+#else
+#define DEVICE
+#endif
+
+#include <float.h>
+#include <math.h>
+#include <stdio.h>
+#include <numpy/npy_math.h>
+
+
+#define MINLOG     -170.0
+#define MAXLOG     +170.0
+#define MAXGAM     171.624376956302725
+#define EPSILON     2.2204460492503131e-16
+
+DEVICE static double pseries(double, double, double);
+DEVICE static double incbcf(double, double, double);
+DEVICE static double incbd(double, double, double);
+
+static double big = 4.503599627370496e15;
+static double biginv =  2.22044604925031308085e-16;
+
+
+DEVICE double BetaInc(double a, double b, double x)
+{
+    double xc, y, w, t;
+    /* check function arguments */
+    if (a <= 0.0) return NPY_NAN;
+    if (b <= 0.0) return NPY_NAN;
+    if (x < 0.0) return NPY_NAN;
+    if (1.0 < x) return NPY_NAN;
+
+    /* some special cases */
+    if (x == 0.0) return 0.0;
+    if (x == 1.0) return 1.0;
+
+    if ( (b * x) <= 1.0 && x <= 0.95)
+    {
+        return pseries(a, b, x);
+    }
+
+    xc = 1.0 - x;
+    /* reverse a and b if x is greater than the mean */
+    if (x > (a / (a + b)))
+    {
+        t = BetaInc(b, a, xc);
+        if (t <= EPSILON) return 1.0 - EPSILON;
+        return 1.0 - t;
+    }
+
+    /* Choose expansion for better convergence. */
+    y = x * (a+b-2.0) - (a-1.0);
+    if( y < 0.0 )
+        w = incbcf( a, b, x );
+    else
+        w = incbd( a, b, x ) / xc;
+
+    y = a * log(x);
+    t = b * log(xc);
+    if( (a+b) < MAXGAM && fabs(y) < MAXLOG && fabs(t) < MAXLOG )
+    {
+        t = pow(xc, b);
+        t *= pow(x, a);
+        t /= a;
+        t *= w;
+        t *= tgamma(a + b) / (tgamma(a) * tgamma(b));
+
+        return t;
+    }
+
+    /* Resort to logarithms.  */
+    y += t + lgamma(a+b) - lgamma(a) - lgamma(b);
+    y += log(w / a);
+    if( y < MINLOG )
+        t = 0.0;
+    else
+        t = exp(y);
+
+    return t;
+}
+
+/* Continued fraction expansion #1
+ * for incomplete beta integral
+ */
+
+DEVICE static double incbcf(double a, double b, double x)
+{
+    double xk, pk, pkm1, pkm2, qk, qkm1, qkm2;
+    double k1, k2, k3, k4, k5, k6, k7, k8;
+    double r, t, ans, thresh;
+    int n;
+
+    k1 = a;
+    k2 = a + b;
+    k3 = a;
+    k4 = a + 1.0;
+    k5 = 1.0;
+    k6 = b - 1.0;
+    k7 = k4;
+    k8 = a + 2.0;
+
+    pkm2 = 0.0;
+    qkm2 = 1.0;
+    pkm1 = 1.0;
+    qkm1 = 1.0;
+    ans = 1.0;
+    r = 1.0;
+    n = 0;
+    thresh = 3.0 * EPSILON;
+    do
+    {
+
+        xk = -( x * k1 * k2 ) / ( k3 * k4 );
+        pk = pkm1 +  pkm2 * xk;
+        qk = qkm1 +  qkm2 * xk;
+        pkm2 = pkm1;
+        pkm1 = pk;
+        qkm2 = qkm1;
+        qkm1 = qk;
+
+        xk = ( x * k5 * k6 ) / ( k7 * k8 );
+        pk = pkm1 +  pkm2 * xk;
+        qk = qkm1 +  qkm2 * xk;
+        pkm2 = pkm1;
+        pkm1 = pk;
+        qkm2 = qkm1;
+        qkm1 = qk;
+
+        if( qk != 0.0 )
+            r = pk/qk;
+        if( r != 0.0 )
+        {
+            t = fabs( (ans - r) / r );
+            ans = r;
+        }
+        else
+            t = 1.0;
+
+        if( t < thresh )
+            break;
+
+        k1 += 1.0;
+        k2 += 1.0;
+        k3 += 2.0;
+        k4 += 2.0;
+        k5 += 1.0;
+        k6 -= 1.0;
+        k7 += 2.0;
+        k8 += 2.0;
+
+        if( (fabs(qk) + fabs(pk)) > big )
+        {
+            pkm2 *= biginv;
+            pkm1 *= biginv;
+            qkm2 *= biginv;
+            qkm1 *= biginv;
+        }
+        if( (fabs(qk) < biginv) || (fabs(pk) < biginv) )
+        {
+            pkm2 *= big;
+            pkm1 *= big;
+            qkm2 *= big;
+            qkm1 *= big;
+        }
+    }
+    while( ++n < 300 );
+
+    return ans;
+}
+
+/* Continued fraction expansion #2
+ * for incomplete beta integral
+ */
+
+DEVICE static double incbd(double a, double b, double x)
+{
+    double xk, pk, pkm1, pkm2, qk, qkm1, qkm2;
+    double k1, k2, k3, k4, k5, k6, k7, k8;
+    double r, t, ans, z, thresh;
+    int n;
+
+    k1 = a;
+    k2 = b - 1.0;
+    k3 = a;
+    k4 = a + 1.0;
+    k5 = 1.0;
+    k6 = a + b;
+    k7 = a + 1.0;;
+    k8 = a + 2.0;
+
+    pkm2 = 0.0;
+    qkm2 = 1.0;
+    pkm1 = 1.0;
+    qkm1 = 1.0;
+    z = x / (1.0-x);
+    ans = 1.0;
+    r = 1.0;
+    n = 0;
+    thresh = 3.0 * EPSILON;
+    do
+    {
+
+        xk = -( z * k1 * k2 ) / ( k3 * k4 );
+        pk = pkm1 +  pkm2 * xk;
+        qk = qkm1 +  qkm2 * xk;
+        pkm2 = pkm1;
+        pkm1 = pk;
+        qkm2 = qkm1;
+        qkm1 = qk;
+
+        xk = ( z * k5 * k6 ) / ( k7 * k8 );
+        pk = pkm1 +  pkm2 * xk;
+        qk = qkm1 +  qkm2 * xk;
+        pkm2 = pkm1;
+        pkm1 = pk;
+        qkm2 = qkm1;
+        qkm1 = qk;
+
+        if( qk != 0 )
+            r = pk/qk;
+        if( r != 0 )
+        {
+            t = fabs( (ans - r) / r );
+            ans = r;
+        }
+        else
+            t = 1.0;
+
+        if( t < thresh )
+            break;
+
+        k1 += 1.0;
+        k2 -= 1.0;
+        k3 += 2.0;
+        k4 += 2.0;
+        k5 += 1.0;
+        k6 += 1.0;
+        k7 += 2.0;
+        k8 += 2.0;
+
+        if( (fabs(qk) + fabs(pk)) > big )
+        {
+            pkm2 *= biginv;
+            pkm1 *= biginv;
+            qkm2 *= biginv;
+            qkm1 *= biginv;
+        }
+        if( (fabs(qk) < biginv) || (fabs(pk) < biginv) )
+        {
+            pkm2 *= big;
+            pkm1 *= big;
+            qkm2 *= big;
+            qkm1 *= big;
+        }
+    }
+    while( ++n < 300 );
+
+    return ans;
+}
+
+
+/* Power series for incomplete beta integral.
+   Use when b*x is small and x not too close to 1.  */
+
+DEVICE static double pseries(double a, double b, double x)
+{
+    double s, t, u, v, n, t1, z, ai;
+
+    ai = 1.0 / a;
+    u = (1.0 - b) * x;
+    v = u / (a + 1.0);
+    t1 = v;
+    t = u;
+    n = 2.0;
+    s = 0.0;
+    z = EPSILON * ai;
+    while( fabs(v) > z )
+    {
+    u = (n - b) * x / n;
+    t *= u;
+    v = t / (a + n);
+    s += v;
+    n += 1.0;
+    }
+    s += t1;
+    s += ai;
+
+    u = a * log(x);
+    if( (a+b) < MAXGAM && fabs(u) < MAXLOG )
+    {
+    t = tgamma(a + b) / (tgamma(a) * tgamma(b));
+    s = s * t * pow(x,a);
+    }
+    else
+    {
+    t = lgamma(a + b) - lgamma(a) - lgamma(b) + u + log(s);
+    if( t < MINLOG )
+    s = 0.0;
+    else
+    s = exp(t);
+    }
+    return s;
+}

{pytensor-2.24.1 → pytensor-2.24.2}/pytensor/scalar/math.py

@@ -1495,8 +1495,25 @@ class BetaInc(ScalarOp):
             ),
         ]
 
-    def c_code(self, *args, **kwargs):
-        raise NotImplementedError()
+    def c_support_code(self, **kwargs):
+        with open(os.path.join(os.path.dirname(__file__), "c_code", "incbet.c")) as f:
+            raw = f.read()
+            return raw
+
+    def c_code(self, node, name, inp, out, sub):
+        (a, b, x) = inp
+        (z,) = out
+        if (
+            node.inputs[0].type in float_types
+            and node.inputs[1].type in float_types
+            and node.inputs[2].type in float_types
+        ):
+            return f"""{z} = BetaInc({a}, {b}, {x});"""
+
+        raise NotImplementedError("type not supported", type)
+
+    def c_code_cache_version(self):
+        return (1,)
 
 
 betainc = BetaInc(upgrade_to_float_no_complex, name="betainc")

{pytensor-2.24.1 → pytensor-2.24.2}/pytensor/sparse/basic.py

@@ -4316,23 +4316,26 @@ def block_diag(*matrices: TensorVariable, format: Literal["csc", "csr"] = "csc")
     --------
     Create a sparse block diagonal matrix from two sparse 2x2 matrices:
 
-    ..code-block:: python
+    .. testcode::
         import numpy as np
         from pytensor.sparse import block_diag
         from scipy.sparse import csr_matrix
 
         A = csr_matrix([[1, 2], [3, 4]])
         B = csr_matrix([[5, 6], [7, 8]])
-        result_sparse = block_diag(A, B, format='csr', name='X')
+        result_sparse = block_diag(A, B, format='csr')
 
         print(result_sparse)
-        >>>  SparseVariable{csr,int32}
-
         print(result_sparse.toarray().eval())
-        >>> array([[1, 2, 0, 0],
-        >>> [3, 4, 0, 0],
-        >>> [0, 0, 5, 6],
-        >>> [0, 0, 7, 8]])
+
+    .. testoutput::
+
+        SparseVariable{csr,int64}
+        [[1 2 0 0]
+         [3 4 0 0]
+         [0 0 5 6]
+         [0 0 7 8]]
+
     """
     if len(matrices) == 1:
         return matrices