pytensor 2.18.1__tar.gz → 2.18.2__tar.gz

{pytensor-2.18.1 → pytensor-2.18.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pytensor
-Version: 2.18.1
+Version: 2.18.2
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License: .. _license:

{pytensor-2.18.1 → pytensor-2.18.2}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2023-11-24T14:59:04+0000",
+ "date": "2023-12-10T11:56:35+0000",
  "dirty": false,
  "error": null,
- "full-revisionid": "7ecb9f8c6b6a2eac940947bac955a10785240667",
- "version": "2.18.1"
+ "full-revisionid": "68b41a480bc9fed69db7331f60543e94241adbd5",
+ "version": "2.18.2"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.18.1 → pytensor-2.18.2}/pytensor/compile/function/types.py

@@ -1474,7 +1474,7 @@ class FunctionMaker:
         if getattr(mode, "profile", None):
             raise TypeError("profile passed via 'mode'. This isn't supported anymore")
         self.profile = profile
-        if profile:
+        if profile and config.cxx:
             # This is very important:
             # 1) We preload the cache here to not have its timing
             #    included with the rewrites.
@@ -1482,7 +1482,12 @@ class FunctionMaker:
             #    too much execution time during testing as we compile
             #    much more functions then the number of compile c
             #    module.
+            start_get_cache = time.perf_counter()
             pytensor.link.c.basic.get_module_cache().refresh()
+            get_cache_time = time.perf_counter() - start_get_cache
+            self.profile.linker_time += get_cache_time
+            self.profile.preload_cache_time += get_cache_time
+
         # Handle the case where inputs and/or outputs is a single
         # Variable (not in a list)
         unpack_single = False
@@ -1722,7 +1727,8 @@ def orig_function(
 
     """
 
-    t1 = time.perf_counter()
+    if profile:
+        t1 = time.perf_counter()
     mode = pytensor.compile.mode.get_mode(mode)
 
     inputs = list(map(convert_function_input, inputs))
@@ -1755,8 +1761,8 @@ def orig_function(
         with config.change_flags(compute_test_value="off"):
             fn = m.create(defaults)
     finally:
-        t2 = time.perf_counter()
-        if fn and profile:
+        if profile and fn:
+            t2 = time.perf_counter()
             profile.compile_time += t2 - t1
             # TODO: append
             profile.nb_nodes = len(fn.maker.fgraph.apply_nodes)

{pytensor-2.18.1 → pytensor-2.18.2}/pytensor/compile/profiling.py

@@ -269,6 +269,9 @@ class ProfileStats:
     linker_time: float = 0.0
     # time spent linking graph (FunctionMaker.create)
 
+    preload_cache_time: float = 0.0
+    # time spent preloading the cache, so it does not affect rewrites profiling
+
     import_time: float = 0.0
     # time spent in importing compiled python module.
 
@@ -811,6 +814,7 @@ class ProfileStats:
             ),
             file=file,
         )
+        print(f"       C-cache preloading {self.preload_cache_time:e}s", file=file)
         print(f"       Import time {self.import_time:e}s", file=file)
         print(
             f"       Node make_thunk time {self.linker_node_make_thunks:e}s", file=file

{pytensor-2.18.1 → pytensor-2.18.2}/pytensor/graph/basic.py

@@ -1777,6 +1777,7 @@ def equal_computations(
     ys: list[Union[np.ndarray, Variable]],
     in_xs: Optional[list[Variable]] = None,
     in_ys: Optional[list[Variable]] = None,
+    strict_dtype=True,
 ) -> bool:
     """Checks if PyTensor graphs represent the same computations.
 
@@ -1908,7 +1909,10 @@ def equal_computations(
                         if dx != dy:
                             if isinstance(dx, Constant) and isinstance(dy, Constant):
                                 if not dx.equals(dy):
-                                    return False
+                                    if strict_dtype:
+                                        return False
+                                    elif not np.array_equal(dx.data, dy.data):
+                                        return False
                             else:
                                 return False
 

{pytensor-2.18.1 → pytensor-2.18.2}/pytensor/graph/replace.py

@@ -3,7 +3,13 @@ from collections.abc import Iterable, Mapping, Sequence
 from functools import partial, singledispatch
 from typing import Optional, Union, cast, overload
 
-from pytensor.graph.basic import Apply, Constant, Variable, truncated_graph_inputs
+from pytensor.graph.basic import (
+    Apply,
+    Constant,
+    Variable,
+    io_toposort,
+    truncated_graph_inputs,
+)
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.op import Op
 
@@ -295,19 +301,14 @@ def vectorize_graph(
     inputs = truncated_graph_inputs(seq_outputs, ancestors_to_include=replace.keys())
     new_inputs = [replace.get(inp, inp) for inp in inputs]
 
-    def transform(var: Variable) -> Variable:
-        if var in inputs:
-            return new_inputs[inputs.index(var)]
+    vect_vars = dict(zip(inputs, new_inputs))
+    for node in io_toposort(inputs, seq_outputs):
+        vect_inputs = [vect_vars.get(inp, inp) for inp in node.inputs]
+        vect_node = vectorize_node(node, *vect_inputs)
+        for output, vect_output in zip(node.outputs, vect_node.outputs):
+            vect_vars[output] = vect_output
 
-        node = var.owner
-        batched_inputs = [transform(inp) for inp in node.inputs]
-        batched_node = vectorize_node(node, *batched_inputs)
-        batched_var = batched_node.outputs[var.owner.outputs.index(var)]
-
-        return cast(Variable, batched_var)
-
-    # TODO: MergeOptimization or node caching?
-    seq_vect_outputs = [transform(out) for out in seq_outputs]
+    seq_vect_outputs = [vect_vars[out] for out in seq_outputs]
 
     if isinstance(outputs, Sequence):
         return seq_vect_outputs

{pytensor-2.18.1 → pytensor-2.18.2}/pytensor/link/c/cmodule.py

@@ -2718,6 +2718,7 @@ def default_blas_ldflags():
                     found = True
                     break
             if not found:
+                _logger.debug("Required file '%s' not found", req)
                 raise RuntimeError(f"Required file {req} not found")
         return libs
 
@@ -2779,9 +2780,15 @@ def default_blas_ldflags():
         res = try_blas_flag(flags)
         if res:
             if any("mkl" in flag for flag in flags):
-                check_mkl_openmp()
+                try:
+                    check_mkl_openmp()
+                except Exception as e:
+                    _logger.debug(e)
+            _logger.debug("The following blas flags will be used: '%s'", res)
             return res
         else:
+            _logger.debug(f"Supplied flags {res} failed to compile")
+            _logger.debug("Supplied flags '%s' failed to compile", res)
             raise RuntimeError(f"Supplied flags {flags} failed to compile")
 
     # If no compiler is available we default to empty ldflags
@@ -2802,6 +2809,11 @@ def default_blas_ldflags():
         # directory. We will include both in our searched library dirs
         searched_library_dirs.append(os.path.join(sys.prefix, "Library", "bin"))
         searched_library_dirs.append(os.path.join(sys.prefix, "Library", "lib"))
+    search_dirs = "\n".join(searched_library_dirs)
+    _logger.debug(
+        "Will search for BLAS libraries in the following directories:\n%s",
+        search_dirs,
+    )
     all_libs = [
         l
         for path in [
@@ -2817,6 +2829,7 @@ def default_blas_ldflags():
         maybe_add_to_os_environ_pathlist("PATH", rpath)
     try:
         # 1. Try to use MKL with INTEL OpenMP threading
+        _logger.debug("Checking MKL flags with intel threading")
         return check_libs(
             all_libs,
             required_libs=[
@@ -2829,19 +2842,21 @@ def default_blas_ldflags():
             extra_compile_flags=[f"-Wl,-rpath,{rpath}"] if rpath is not None else [],
             cxx_library_dirs=cxx_library_dirs,
         )
-    except Exception:
-        pass
+    except Exception as e:
+        _logger.debug(e)
     try:
         # 2. Try to use MKL with GNU OpenMP threading
+        _logger.debug("Checking MKL flags with GNU OpenMP threading")
         return check_libs(
             all_libs,
             required_libs=["mkl_core", "mkl_rt", "mkl_gnu_thread", "gomp", "pthread"],
             extra_compile_flags=[f"-Wl,-rpath,{rpath}"] if rpath is not None else [],
             cxx_library_dirs=cxx_library_dirs,
         )
-    except Exception:
-        pass
+    except Exception as e:
+        _logger.debug(e)
     try:
+        _logger.debug("Checking Lapack + blas")
         # 3. Try to use LAPACK + BLAS
         return check_libs(
             all_libs,
@@ -2849,20 +2864,22 @@ def default_blas_ldflags():
             extra_compile_flags=[f"-Wl,-rpath,{rpath}"] if rpath is not None else [],
             cxx_library_dirs=cxx_library_dirs,
         )
-    except Exception:
-        pass
+    except Exception as e:
+        _logger.debug(e)
     try:
         # 4. Try to use BLAS alone
+        _logger.debug("Checking blas alone")
         return check_libs(
             all_libs,
             required_libs=["blas", "cblas"],
             extra_compile_flags=[f"-Wl,-rpath,{rpath}"] if rpath is not None else [],
             cxx_library_dirs=cxx_library_dirs,
         )
-    except Exception:
-        pass
+    except Exception as e:
+        _logger.debug(e)
     try:
         # 5. Try to use openblas
+        _logger.debug("Checking openblas")
         return check_libs(
             all_libs,
             required_libs=["openblas", "gfortran", "gomp", "m"],
@@ -2871,8 +2888,9 @@ def default_blas_ldflags():
             else ["-fopenmp"],
             cxx_library_dirs=cxx_library_dirs,
         )
-    except Exception:
-        pass
+    except Exception as e:
+        _logger.debug(e)
+    _logger.debug("Failed to identify blas ldflags. Will leave them empty.")
     return ""
 
 

{pytensor-2.18.1 → pytensor-2.18.2}/pytensor/link/jax/dispatch/nlinalg.py

@@ -99,9 +99,7 @@ def jax_funcify_BatchedDot(op, **kwargs):
     def batched_dot(a, b):
         if a.shape[0] != b.shape[0]:
             raise TypeError("Shapes must match in the 0-th dimension")
-        if a.ndim == 2 or b.ndim == 2:
-            return jnp.einsum("n...j,nj...->n...", a, b)
-        return jnp.einsum("nij,njk->nik", a, b)
+        return jnp.matmul(a, b)
 
     return batched_dot
 

{pytensor-2.18.1 → pytensor-2.18.2}/pytensor/link/numba/dispatch/basic.py

@@ -895,6 +895,8 @@ def numba_funcify_BatchedDot(op, node, **kwargs):
 
     @numba_njit
     def batched_dot(x, y):
+        # Numba does not support 3D matmul
+        # https://github.com/numba/numba/issues/3804
         shape = x.shape[:-1] + y.shape[2:]
         z0 = np.empty(shape, dtype=dtype)
         for i in range(z0.shape[0]):

{pytensor-2.18.1 → pytensor-2.18.2}/pytensor/tensor/basic.py

@@ -22,10 +22,12 @@ import pytensor.scalar.sharedvar
 from pytensor import compile, config, printing
 from pytensor import scalar as aes
 from pytensor.gradient import DisconnectedType, grad_undefined
+from pytensor.graph import RewriteDatabaseQuery
 from pytensor.graph.basic import Apply, Constant, Variable
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.op import Op
-from pytensor.graph.rewriting.utils import rewrite_graph
+from pytensor.graph.replace import _vectorize_node
+from pytensor.graph.rewriting.db import EquilibriumDB
 from pytensor.graph.type import HasShape, Type
 from pytensor.link.c.op import COp
 from pytensor.link.c.params_type import ParamsType
@@ -40,6 +42,7 @@ from pytensor.tensor import (
     as_tensor_variable,
     get_vector_length,
 )
+from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.elemwise import DimShuffle, Elemwise, scalar_elemwise
 from pytensor.tensor.exceptions import NotScalarConstantError
 from pytensor.tensor.shape import (
@@ -1356,6 +1359,45 @@ def identity_like(x, dtype: Optional[Union[str, np.generic, np.dtype]] = None):
     return eye(_x.shape[0], _x.shape[1], k=0, dtype=dtype)
 
 
+class CachedEquilibrimDB(EquilibriumDB):
+    """A subclass of EquilibriumDB that allows caching of a default query for faster reuse."""
+
+    def __init__(self, default_query):
+        super().__init__()
+        self._default_query = default_query
+        self._cached_default_query = None
+
+    def register(self, *args, **kwargs):
+        # If new rewrites are registered, the default cached query is void
+        self.cached_default_query = None
+        super().register(*args, **kwargs)
+
+    @property
+    def default_query(self):
+        if self._cached_default_query is None:
+            self._cached_default_query = self.query(self._default_query)
+        return self._cached_default_query
+
+
+infer_shape_db = CachedEquilibrimDB(
+    default_query=RewriteDatabaseQuery(include=("infer_shape",))
+)
+
+
+def register_infer_shape(rewrite, *tags, **kwargs):
+    if isinstance(rewrite, str):
+
+        def register(inner_lopt):
+            return register_infer_shape(inner_lopt, rewrite, *tags, **kwargs)
+
+        return register
+    else:
+        name = kwargs.pop("name", None) or rewrite.__name__
+
+        infer_shape_db.register(name, rewrite, *tags, "infer_shape", **kwargs)
+        return rewrite
+
+
 def infer_static_shape(
     shape: Union[Variable, Sequence[Union[Variable, int]]]
 ) -> tuple[Sequence["TensorLike"], Sequence[Optional[int]]]:
@@ -1390,14 +1432,16 @@ def infer_static_shape(
 
         raise TypeError(f"Shapes must be scalar integers; got {s_as_str}")
 
-    sh = [check_type(as_tensor_variable(s, ndim=0)) for s in shape]
+    sh = folded_shape = [check_type(as_tensor_variable(s, ndim=0)) for s in shape]
+
+    if not all(isinstance(s, Constant) for s in folded_shape):
+        shape_fg = FunctionGraph(outputs=sh, features=[ShapeFeature()], clone=True)
+        with config.change_flags(optdb__max_use_ratio=10, cxx=""):
+            infer_shape_db.default_query.rewrite(shape_fg)
+            if not all(isinstance(s, Constant) for s in shape_fg.outputs):
+                topo_constant_folding.rewrite(shape_fg)
+        folded_shape = shape_fg.outputs
 
-    shape_fg = FunctionGraph(
-        outputs=sh,
-        features=[ShapeFeature()],
-        clone=True,
-    )
-    folded_shape = rewrite_graph(shape_fg, custom_rewrite=topo_constant_folding).outputs
     static_shape = tuple(
         s.data.item() if isinstance(s, Constant) else None for s in folded_shape
     )
@@ -1615,16 +1659,22 @@ class Alloc(COp):
         if not clients:
             return False
 
-        for client in clients:
-            if client[0] == "output":
+        for client, idx in clients:
+            if client == "output":
                 # If the output is a constant, it will have to be deepcopied
                 # each time the function is called.  So we do not fold.
                 return False
+            # Allow alloc to be lifted out of Elemwise before constant folding it
+            elif isinstance(client.op, Elemwise):
+                return None
+            # Same for Blockwise, unless it has no batch_dims
+            elif isinstance(client.op, Blockwise) and client.op.batch_ndim(client):
+                return None
             elif (
                 # The following ops work inplace of their input id 0.
-                client[1] == 0
+                idx == 0
                 and isinstance(
-                    client[0].op,
+                    client.op,
                     (
                         # Ops that will work inplace on the Alloc. So if they
                         # get constant_folded, they would copy the
@@ -3455,10 +3505,17 @@ class ExtractDiag(Op):
 
         if x.ndim < 2:
             raise ValueError("ExtractDiag needs an input with 2 or more dimensions", x)
+
+        out_shape = [
+            st_dim
+            for i, st_dim in enumerate(x.type.shape)
+            if i not in (self.axis1, self.axis2)
+        ] + [None]
+
         return Apply(
             self,
             [x],
-            [x.type.clone(dtype=x.dtype, shape=(None,) * (x.ndim - 1))()],
+            [x.type.clone(dtype=x.dtype, shape=tuple(out_shape))()],
         )
 
     def perform(self, node, inputs, outputs):
@@ -3559,6 +3616,17 @@ def diagonal(a, offset=0, axis1=0, axis2=1):
     return ExtractDiag(offset, axis1, axis2)(a)
 
 
+@_vectorize_node.register(ExtractDiag)
+def vectorize_extract_diag(op: ExtractDiag, node, batched_x):
+    batched_ndims = batched_x.type.ndim - node.inputs[0].type.ndim
+    return diagonal(
+        batched_x,
+        offset=op.offset,
+        axis1=op.axis1 + batched_ndims,
+        axis2=op.axis2 + batched_ndims,
+    ).owner
+
+
 def trace(a, offset=0, axis1=0, axis2=1):
     """
     Returns the sum along diagonals of the array.