pytensor 2.31.6__tar.gz → 2.31.7__tar.gz

{pytensor-2.31.6/pytensor.egg-info → pytensor-2.31.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 2.31.6
+Version: 2.31.7
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License-Expression: BSD-3-Clause

{pytensor-2.31.6 → pytensor-2.31.7}/doc/conf.py

@@ -1,6 +1,7 @@
 import os
 import inspect
 import sys
+
 import pytensor
 from pathlib import Path
 
@@ -12,6 +13,7 @@ extensions = [
     "sphinx.ext.autodoc",
     "sphinx.ext.todo",
     "sphinx.ext.doctest",
+    "sphinx_copybutton",
     "sphinx.ext.napoleon",
     "sphinx.ext.linkcode",
     "sphinx.ext.mathjax",
@@ -86,8 +88,7 @@ today_fmt = "%B %d, %Y"
 
 # List of directories, relative to source directories, that shouldn't be
 # searched for source files.
-exclude_dirs = ["images", "scripts", "sandbox"]
-exclude_patterns = ['page_footer.md', '**/*.myst.md']
+exclude_patterns = ["README.md", "images/*", "page_footer.md", "**/*.myst.md"]
 
 # The reST default role (used for this markup: `text`) to use for all
 # documents.
@@ -235,24 +236,41 @@ htmlhelp_basename = "pytensor_doc"
 # Resolve function
 # This function is used to populate the (source) links in the API
 def linkcode_resolve(domain, info):
-    def find_source():
+    def find_obj() -> object:
         # try to find the file and line number, based on code from numpy:
         # https://github.com/numpy/numpy/blob/master/doc/source/conf.py#L286
         obj = sys.modules[info["module"]]
         for part in info["fullname"].split("."):
             obj = getattr(obj, part)
+        return obj
 
+    def find_source(obj):
         fn = Path(inspect.getsourcefile(obj))
-        fn = fn.relative_to(Path(__file__).parent)
+        fn = fn.relative_to(Path(pytensor.__file__).parent)
         source, lineno = inspect.getsourcelines(obj)
         return fn, lineno, lineno + len(source) - 1
 
+    def fallback_source():
+        return info["module"].replace(".", "/") + ".py"
+
     if domain != "py" or not info["module"]:
         return None
+
     try:
-        filename = "pytensor/%s#L%d-L%d" % find_source()
+        obj = find_obj()
     except Exception:
-        filename = info["module"].replace(".", "/") + ".py"
+        filename = fallback_source()
+    else:
+        try:
+            filename = "pytensor/%s#L%d-L%d" % find_source(obj)
+        except Exception:
+            # warnings.warn(f"Could not find source code for {domain}:{info}")
+            try:
+                filename = obj.__module__.replace(".", "/") + ".py"
+            except AttributeError:
+                # Some objects do not have a __module__ attribute (?)
+                filename = fallback_source()
+
     import subprocess
 
     tag = subprocess.Popen(

{pytensor-2.31.6 → pytensor-2.31.7}/doc/environment.yml

@@ -13,7 +13,9 @@ dependencies:
   - mock
   - pillow
   - pymc-sphinx-theme
+  - sphinx-copybutton
   - sphinx-design
+  - sphinx-sitemap
   - pygments
   - pydot
   - ipython
@@ -23,5 +25,4 @@ dependencies:
   - ablog
   - pip
   - pip:
-    - sphinx_sitemap
     - -e ..

{pytensor-2.31.6 → pytensor-2.31.7}/doc/library/index.rst

@@ -20,14 +20,12 @@ Modules
    d3viz/index
    graph/index
    gradient
-   misc/pkl_utils
    printing
-   scalar/index
    scan
    sparse/index
-   sparse/sandbox
    tensor/index
    typed_list
+   xtensor/index
 
 .. module:: pytensor
    :platform: Unix, Windows

pytensor-2.31.7/doc/library/xtensor/index.md

@@ -0,0 +1,101 @@
+(libdoc_xtensor)=
+# `xtensor` -- XTensor operations
+
+This module implements as abstraction layer on regular tensor operations, that behaves like Xarray.
+
+A new type {class}`pytensor.xtensor.type.XTensorType`, generalizes the {class}`pytensor.tensor.TensorType` 
+with the addition of a `dims` attribute, that labels the dimensions of the tensor. 
+
+Variables of XTensorType (i.e.,  {class}`pytensor.xtensor.type.XTensorVariable`s) are the symbolic counterpart
+to xarray DataArray objects.
+
+The module implements several PyTensor operations {class}`pytensor.xtensor.basic.XOp`s, whose signature mimics that of 
+xarray (and xarray_einstats) DataArray operations. These operations, unlike most regular PyTensor operations, cannot 
+be directly evaluated, but require a rewrite (lowering) into a regular tensor graph that can itself be evaluated as usual.
+
+Like regular PyTensor, we don't need an Op for every possible method or function in the public API of xarray.
+If the existing XOps can be composed to produce the desired result, then we can use them directly.
+
+## Coordinates
+For now, there's no analogous of xarray coordinates, so you won't be able to do coordinate operations like `.sel`.
+The graphs produced by an xarray program without coords are much more amenable to the numpy-like backend of PyTensor.
+Coords involve aspects of Pandas/database query and joining that are not trivially expressible in PyTensor.
+
+## Example
+
+
+```{testcode}
+
+import pytensor.tensor as pt
+import pytensor.xtensor as ptx
+
+a = pt.tensor("a", shape=(3,))
+b = pt.tensor("b", shape=(4,))
+
+ax = ptx.as_xtensor(a, dims=["x"])
+bx = ptx.as_xtensor(b, dims=["y"])
+
+zx = ax + bx
+assert zx.type == ptx.type.XTensorType("float64", dims=["x", "y"], shape=(3, 4))
+
+z = zx.values
+z.dprint()
+```
+
+
+```{testoutput}
+
+TensorFromXTensor [id A]
+ └─ XElemwise{scalar_op=Add()} [id B]
+    ├─ XTensorFromTensor{dims=('x',)} [id C]
+    │  └─ a [id D]
+    └─ XTensorFromTensor{dims=('y',)} [id E]
+       └─ b [id F]
+```
+
+Once we compile the graph, no XOps are left.
+
+```{testcode}
+
+import pytensor
+
+with pytensor.config.change_flags(optimizer_verbose=True):
+    fn = pytensor.function([a, b], z)
+
+```
+
+```{testoutput}
+
+rewriting: rewrite lower_elemwise replaces XElemwise{scalar_op=Add()}.0 of XElemwise{scalar_op=Add()}(XTensorFromTensor{dims=('x',)}.0, XTensorFromTensor{dims=('y',)}.0) with XTensorFromTensor{dims=('x', 'y')}.0 of XTensorFromTensor{dims=('x', 'y')}(Add.0)
+rewriting: rewrite useless_tensor_from_xtensor replaces TensorFromXTensor.0 of TensorFromXTensor(XTensorFromTensor{dims=('x',)}.0) with a of None
+rewriting: rewrite useless_tensor_from_xtensor replaces TensorFromXTensor.0 of TensorFromXTensor(XTensorFromTensor{dims=('y',)}.0) with b of None
+rewriting: rewrite useless_tensor_from_xtensor replaces TensorFromXTensor.0 of TensorFromXTensor(XTensorFromTensor{dims=('x', 'y')}.0) with Add.0 of Add(ExpandDims{axis=1}.0, ExpandDims{axis=0}.0)
+
+```
+
+```{testcode}
+
+fn.dprint()
+```
+
+```{testoutput}
+
+Add [id A] 2
+ ├─ ExpandDims{axis=1} [id B] 1
+ │  └─ a [id C]
+ └─ ExpandDims{axis=0} [id D] 0
+    └─ b [id E]
+```
+
+
+## Index
+
+:::{toctree}
+:maxdepth: 1
+
+module_functions
+math
+linalg
+random
+type
+:::

pytensor-2.31.7/doc/library/xtensor/linalg.md

@@ -0,0 +1,7 @@
+(libdoc_xtensor_linalg)=
+# `xtensor.linalg` -- Linear algebra operations
+
+```{eval-rst}
+.. automodule:: pytensor.xtensor.linalg
+   :members:
+```

pytensor-2.31.7/doc/library/xtensor/math.md

@@ -0,0 +1,8 @@
+(libdoc_xtensor_math)=
+# `xtensor.math` Mathematical operations
+
+```{eval-rst}
+.. automodule:: pytensor.xtensor.math
+   :members:
+   :exclude-members: XDot, dot
+```

pytensor-2.31.7/doc/library/xtensor/module_functions.md

@@ -0,0 +1,7 @@
+(libdoc_xtensor_module_function)=
+# `xtensor` -- Module level operations
+
+```{eval-rst}
+.. automodule:: pytensor.xtensor
+   :members: broadcast, concat, dot, full_like, ones_like, zeros_like
+```

pytensor-2.31.7/doc/library/xtensor/random.md

@@ -0,0 +1,7 @@
+(libdoc_xtensor_random)=
+# `xtensor.random` Random number generator operations
+
+```{eval-rst}
+.. automodule:: pytensor.xtensor.random
+   :members:
+```

pytensor-2.31.7/doc/library/xtensor/type.md

@@ -0,0 +1,21 @@
+(libdoc_xtenor_type)=
+
+# `xtensor.type` -- Types and Variables
+
+## XTensorVariable creation functions
+
+```{eval-rst}
+.. automodule:: pytensor.xtensor.type
+   :members: xtensor, xtensor_constant, as_xtensor
+
+```
+
+## XTensor Type and Variable classes
+
+```{eval-rst}
+.. automodule:: pytensor.xtensor.type
+   :noindex:
+   :members: XTensorType, XTensorVariable, XTensorConstant
+```
+
+

{pytensor-2.31.6 → pytensor-2.31.7}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-07-02T21:17:39+0200",
+ "date": "2025-07-08T17:42:30+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "815671d5f53ab7bf0cabe9ba685b70285da38a27",
- "version": "2.31.6"
+ "full-revisionid": "4ce092fe4df857a2bf88a6a367b4a1592e10c0ee",
+ "version": "2.31.7"
 }
 '''  # END VERSION_JSON
 

pytensor-2.31.7/pytensor/link/jax/dispatch/blockwise.py

@@ -0,0 +1,21 @@
+import jax.numpy as jnp
+
+from pytensor.link.jax.dispatch import jax_funcify
+from pytensor.tensor.blockwise import Blockwise
+
+
+@jax_funcify.register(Blockwise)
+def jax_funcify_Blockwise(op: Blockwise, node, **kwargs):
+    signature = op.signature
+    core_node = op._create_dummy_core_node(
+        node.inputs, propagate_unbatched_core_inputs=True
+    )
+    core_fn = jax_funcify(core_node.op, node=core_node, **kwargs)
+
+    vect_fn = jnp.vectorize(core_fn, signature=signature)
+
+    def blockwise_fn(*inputs):
+        op._check_runtime_broadcast(node, inputs)
+        return vect_fn(*inputs)
+
+    return blockwise_fn

pytensor-2.31.7/pytensor/link/jax/dispatch/signal/conv.py

@@ -0,0 +1,24 @@
+import jax
+
+from pytensor.link.jax.dispatch import jax_funcify
+from pytensor.tensor.basic import get_underlying_scalar_constant_value
+from pytensor.tensor.exceptions import NotScalarConstantError
+from pytensor.tensor.signal.conv import Convolve1d
+
+
+@jax_funcify.register(Convolve1d)
+def jax_funcify_Convolve1d(op, node, **kwargs):
+    _, _, full_mode = node.inputs
+    try:
+        full_mode = get_underlying_scalar_constant_value(full_mode)
+    except NotScalarConstantError:
+        raise NotImplementedError(
+            "Cannot compile Convolve1D to jax without static mode"
+        )
+    static_mode = "full" if full_mode else "valid"
+
+    def conv1d(data, kernel, _runtime_full_mode):
+        # _runtime_full_mode is not used, as we only support static mode
+        return jax.numpy.convolve(data, kernel, mode=static_mode)
+
+    return conv1d

{pytensor-2.31.6 → pytensor-2.31.7}/pytensor/link/numba/dispatch/blockwise.py

@@ -16,7 +16,7 @@ from pytensor.tensor import TensorVariable, get_vector_length
 from pytensor.tensor.blockwise import Blockwise, BlockwiseWithCoreShape
 
 
-@numba_funcify.register
+@numba_funcify.register(BlockwiseWithCoreShape)
 def numba_funcify_Blockwise(op: BlockwiseWithCoreShape, node, **kwargs):
     [blockwise_node] = op.fgraph.apply_nodes
     blockwise_op: Blockwise = blockwise_node.op
@@ -26,7 +26,8 @@ def numba_funcify_Blockwise(op: BlockwiseWithCoreShape, node, **kwargs):
     core_shapes_len = tuple(get_vector_length(sh) for sh in node.inputs[nin:])
 
     core_node = blockwise_op._create_dummy_core_node(
-        cast(tuple[TensorVariable], blockwise_node.inputs)
+        cast(tuple[TensorVariable], node.inputs[:nin]),
+        propagate_unbatched_core_inputs=True,
     )
     core_op_fn = numba_funcify(
         core_op,

pytensor-2.31.7/pytensor/link/numba/dispatch/signal/conv.py

@@ -0,0 +1,69 @@
+import numpy as np
+from numba.np.arraymath import _get_inner_prod
+
+from pytensor.link.numba.dispatch import numba_funcify
+from pytensor.link.numba.dispatch.basic import numba_njit
+from pytensor.tensor.signal.conv import Convolve1d
+
+
+@numba_funcify.register(Convolve1d)
+def numba_funcify_Convolve1d(op, node, **kwargs):
+    # This specialized version is faster than the overloaded numba np.convolve
+    a_dtype, b_dtype = node.inputs[0].type.dtype, node.inputs[1].type.dtype
+    out_dtype = node.outputs[0].type.dtype
+    innerprod = _get_inner_prod(a_dtype, b_dtype)
+
+    @numba_njit
+    def valid_convolve1d(x, y):
+        nx = len(x)
+        ny = len(y)
+        if nx < ny:
+            x, y = y, x
+            nx, ny = ny, nx
+        y_flipped = y[::-1]
+
+        length = nx - ny + 1
+        ret = np.empty(length, out_dtype)
+
+        for i in range(length):
+            ret[i] = innerprod(x[i : i + ny], y_flipped)
+
+        return ret
+
+    @numba_njit
+    def full_convolve1d(x, y):
+        nx = len(x)
+        ny = len(y)
+        if nx < ny:
+            x, y = y, x
+            nx, ny = ny, nx
+        y_flipped = y[::-1]
+
+        length = nx + ny - 1
+        ret = np.empty(length, out_dtype)
+        idx = 0
+
+        for i in range(ny - 1):
+            k = i + 1
+            ret[idx] = innerprod(x[:k], y_flipped[-k:])
+            idx = idx + 1
+
+        for i in range(nx - ny + 1):
+            ret[idx] = innerprod(x[i : i + ny], y_flipped)
+            idx = idx + 1
+
+        for i in range(ny - 1):
+            k = ny - i - 1
+            ret[idx] = innerprod(x[-k:], y_flipped[:k])
+            idx = idx + 1
+
+        return ret
+
+    @numba_njit
+    def convolve_1d(x, y, mode):
+        if mode:
+            return full_convolve1d(x, y)
+        else:
+            return valid_convolve1d(x, y)
+
+    return convolve_1d

{pytensor-2.31.6 → pytensor-2.31.7}/pytensor/scalar/math.py

@@ -385,46 +385,56 @@ class Psi(UnaryScalarOp):
             #define DEVICE
             #endif
 
-            #ifndef ga_double
-            #define ga_double double
+            #ifndef M_PI
+            #define M_PI 3.14159265358979323846
             #endif
 
             #ifndef _PSIFUNCDEFINED
             #define _PSIFUNCDEFINED
-            DEVICE double _psi(ga_double x) {
-
-            /*taken from
-            Bernardo, J. M. (1976). Algorithm AS 103:
-            Psi (Digamma) Function. Applied Statistics. 25 (3), 315-317.
-            http://www.uv.es/~bernardo/1976AppStatist.pdf */
+            DEVICE double _psi(double x) {
+
+                /*taken from
+                Bernardo, J. M. (1976). Algorithm AS 103:
+                Psi (Digamma) Function. Applied Statistics. 25 (3), 315-317.
+                http://www.uv.es/~bernardo/1976AppStatist.pdf
+                */
+
+                double y, R, psi_ = 0;
+                double S  = 1.0e-5;
+                double C = 8.5;
+                double S3 = 8.333333333e-2;
+                double S4 = 8.333333333e-3;
+                double S5 = 3.968253968e-3;
+                double D1 = -0.5772156649;
 
-            ga_double y, R, psi_ = 0;
-            ga_double S  = 1.0e-5;
-            ga_double C = 8.5;
-            ga_double S3 = 8.333333333e-2;
-            ga_double S4 = 8.333333333e-3;
-            ga_double S5 = 3.968253968e-3;
-            ga_double D1 = -0.5772156649;
-
-            y = x;
+                if (x <= 0) {
+                    // the digamma function approaches infinity from one side and -infinity from the other, around negative integers and zero
+                    if (x == floor(x)) {
+                        return INFINITY; // note that scipy returns -INF for 0 and NaN for negative integers
+                    }
+
+                    // Use reflection formula
+                    double pi_x = M_PI * x;
+                    double cot_pi_x = cos(pi_x) / sin(pi_x);
+                    return _psi(1.0 - x) - M_PI * cot_pi_x;
+                }
 
-            if (y <= 0.0)
-               return psi_;
+                y = x;
 
-            if (y <= S)
-                return D1 - 1.0/y;
+                if (y <= S)
+                    return D1 - 1.0/y;
 
-            while (y < C) {
-                psi_ = psi_ - 1.0 / y;
-                y = y + 1;
-            }
+                while (y < C) {
+                    psi_ = psi_ - 1.0 / y;
+                    y = y + 1;
+                }
 
-            R = 1.0 / y;
-            psi_ = psi_ + log(y) - .5 * R ;
-            R= R*R;
-            psi_ = psi_ - R * (S3 - R * (S4 - R * S5));
+                R = 1.0 / y;
+                psi_ = psi_ + log(y) - .5 * R ;
+                R= R*R;
+                psi_ = psi_ - R * (S3 - R * (S4 - R * S5));
 
-            return psi_;
+                return psi_;
             }
             #endif
             """
@@ -433,8 +443,8 @@ class Psi(UnaryScalarOp):
         (x,) = inp
         (z,) = out
         if node.inputs[0].type in float_types:
-            return f"""{z} =
-                _psi({x});"""
+            dtype = "npy_" + node.outputs[0].dtype
+            return f"{z} = ({dtype}) _psi({x});"
         raise NotImplementedError("only floating point is implemented")
 
 

{pytensor-2.31.6 → pytensor-2.31.7}/pytensor/tensor/basic.py

@@ -678,10 +678,9 @@ class ScalarFromTensor(COp):
             self, [t], [ps.get_scalar_type(dtype=t.type.dtype).make_variable()]
         )
 
-    def perform(self, node, inp, out_):
-        (s,) = inp
-        (out,) = out_
-        out[0] = s.flatten()[0]
+    def perform(self, node, inputs, output_storage):
+        # not using .item() because that returns a Python scalar, not a numpy scalar
+        output_storage[0][0] = inputs[0][()]
 
     def infer_shape(self, fgraph, node, in_shapes):
         return [()]