pytensor 2.18.0__tar.gz → 2.18.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pytensor
-Version: 2.18.0
+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.0 → pytensor-2.18.2}/doc/extending/using_params.rst

@@ -1,10 +1,10 @@
 .. _extending_op_params:
 
-===============
-Using Op params
-===============
+================
+Using COp params
+================
 
-The Op params is a facility to pass some runtime parameters to the
+The COp params is a facility to pass some runtime parameters to the
 code of an op without modifying it.  It can enable a single instance
 of C code to serve different needs and therefore reduce compilation.
 
@@ -53,7 +53,7 @@ following methods will be used for the type:
   - :meth:`__hash__ <Type.__hash__>`
   - :meth:`values_eq <Type.values_eq>`
 
-Additionally if you want to use your params with C code, you need to extend `COp`
+Additionally, to use your params with C code, you need to extend `COp`
 and implement the following methods:
 
   - :meth:`c_declare <CLinkerType.c_declare>`
@@ -65,24 +65,24 @@ You can also define other convenience methods such as
 :meth:`c_headers <CLinkerType.c_headers>` if you need any special things.
 
 
-Registering the params with your Op
------------------------------------
+Registering the params with your COp
+------------------------------------
 
-To declare that your Op uses params you have to set the class
+To declare that your `COp` uses params you have to set the class
 attribute :attr:`params_type` to an instance of your params Type.
 
 .. note::
 
    If you want to have multiple parameters, PyTensor provides the convenient class
    :class:`pytensor.link.c.params_type.ParamsType` that allows to bundle many parameters into
-   one object that will be available in both Python (as a Python object) and C code (as a struct).
+   one object that will be available to the C code (as a struct).
 
 For example if we decide to use an int as the params the following
 would be appropriate:
 
 .. code-block:: python
 
-   class MyOp(Op):
+   class MyOp(COp):
        params_type = Generic()
 
 After that you need to define a :meth:`get_params` method on your
@@ -115,12 +115,7 @@ Having declared a params for your Op will affect the expected
 signature of :meth:`perform`.  The new expected signature will have an
 extra parameter at the end which corresponds to the params object.
 
-.. warning::
-
-   If you do not account for this extra parameter, the code will fail
-   at runtime if it tries to run the python version.
-
-Also, for the C code, the `sub` dictionary will contain an extra entry
+The `sub` dictionary for `COp`s with params will contain an extra entry
 `'params'` which will map to the variable name of the params object.
 This is true for all methods that receive a `sub` parameter, so this
 means that you can use your params in the :meth:`c_code <COp.c_code>`
@@ -131,7 +126,7 @@ A simple example
 ----------------
 
 This is a simple example which uses a params object to pass a value.
-This `Op` will multiply a scalar input by a fixed floating point value.
+This `COp` will multiply a scalar input by a fixed floating point value.
 
 Since the value in this case is a python float, we chose Generic as
 the params type.
@@ -156,9 +151,10 @@ the params type.
            inp = as_scalar(inp)
            return Apply(self, [inp], [inp.type()])
 
-       def perform(self, node, inputs, output_storage, params):
-           # Here params is a python float so this is ok
-           output_storage[0][0] = inputs[0] * params
+       def perform(self, node, inputs, output_storage):
+           # Because params is a python float we can use `self.mul` directly.
+           # If it's something fancier, call `self.params_type.filter(self.get_params(node))`
+           output_storage[0][0] = inputs[0] * self.mul
 
        def c_code(self, node, name, inputs, outputs, sub):
            return ("%(z)s = %(x)s * PyFloat_AsDouble(%(p)s);" %
@@ -174,7 +170,7 @@ weights.
 
 .. testcode::
 
-   from pytensor.graph.op import Op
+   from pytensor.link.c.op import COp
    from pytensor.link.c.type import Generic
    from pytensor.scalar import as_scalar
 

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

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2023-11-20T14:04:50+0100",
+ "date": "2023-12-10T11:56:35+0000",
  "dirty": false,
  "error": null,
- "full-revisionid": "eb552eef7a8381e47a38ce0694244700693cb978",
- "version": "2.18.0"
+ "full-revisionid": "68b41a480bc9fed69db7331f60543e94241adbd5",
+ "version": "2.18.2"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.18.0 → 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.0 → 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.0 → pytensor-2.18.2}/pytensor/graph/basic.py

@@ -30,7 +30,6 @@ import numpy as np
 from pytensor.configdefaults import config
 from pytensor.graph.utils import (
     MetaObject,
-    MethodNotDefined,
     Scratchpad,
     TestValueError,
     ValidatingScratchpad,
@@ -151,16 +150,6 @@ class Apply(Node, Generic[OpType]):
                     f"The 'outputs' argument to Apply must contain Variable instances with no owner, not {output}"
                 )
 
-    def run_params(self):
-        """
-        Returns the params for the node, or NoParams if no params is set.
-
-        """
-        try:
-            return self.op.get_params(self)
-        except MethodNotDefined:
-            return NoParams
-
     def __getstate__(self):
         d = self.__dict__
         # ufunc don't pickle/unpickle well
@@ -1788,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.
 
@@ -1919,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.0 → pytensor-2.18.2}/pytensor/graph/op.py

@@ -16,15 +16,13 @@ from typing import (
 
 import pytensor
 from pytensor.configdefaults import config
-from pytensor.graph.basic import Apply, NoParams, Variable
+from pytensor.graph.basic import Apply, Variable
 from pytensor.graph.utils import (
     MetaObject,
-    MethodNotDefined,
     TestValueError,
     add_tag_trace,
     get_variable_trace_string,
 )
-from pytensor.link.c.params_type import Params, ParamsType
 
 
 if TYPE_CHECKING:
@@ -37,10 +35,7 @@ StorageMapType = dict[Variable, StorageCellType]
 ComputeMapType = dict[Variable, list[bool]]
 InputStorageType = list[StorageCellType]
 OutputStorageType = list[StorageCellType]
-ParamsInputType = Optional[tuple[Any, ...]]
-PerformMethodType = Callable[
-    [Apply, list[Any], OutputStorageType, ParamsInputType], None
-]
+PerformMethodType = Callable[[Apply, list[Any], OutputStorageType], None]
 BasicThunkType = Callable[[], None]
 ThunkCallableType = Callable[
     [PerformMethodType, StorageMapType, ComputeMapType, Apply], None
@@ -202,7 +197,6 @@ class Op(MetaObject):
 
     itypes: Optional[Sequence["Type"]] = None
     otypes: Optional[Sequence["Type"]] = None
-    params_type: Optional[ParamsType] = None
 
     _output_type_depends_on_input_value = False
     """
@@ -426,7 +420,6 @@ class Op(MetaObject):
         node: Apply,
         inputs: Sequence[Any],
         output_storage: OutputStorageType,
-        params: ParamsInputType = None,
     ) -> None:
         """Calculate the function on the inputs and put the variables in the output storage.
 
@@ -442,8 +435,6 @@ class Op(MetaObject):
             these lists).  Each sub-list corresponds to value of each
             `Variable` in :attr:`node.outputs`.  The primary purpose of this method
             is to set the values of these sub-lists.
-        params
-            A tuple containing the values of each entry in :attr:`Op.__props__`.
 
         Notes
         -----
@@ -481,22 +472,6 @@ class Op(MetaObject):
         """
         return True
 
-    def get_params(self, node: Apply) -> Params:
-        """Try to get parameters for the `Op` when :attr:`Op.params_type` is set to a `ParamsType`."""
-        if isinstance(self.params_type, ParamsType):
-            wrapper = self.params_type
-            if not all(hasattr(self, field) for field in wrapper.fields):
-                # Let's print missing attributes for debugging.
-                not_found = tuple(
-                    field for field in wrapper.fields if not hasattr(self, field)
-                )
-                raise AttributeError(
-                    f"{type(self).__name__}: missing attributes {not_found} for ParamsType."
-                )
-            # ParamsType.get_params() will apply filtering to attributes.
-            return self.params_type.get_params(self)
-        raise MethodNotDefined("get_params")
-
     def prepare_node(
         self,
         node: Apply,
@@ -538,34 +513,12 @@ class Op(MetaObject):
         else:
             p = node.op.perform
 
-        params = node.run_params()
-
-        if params is NoParams:
-            # default arguments are stored in the closure of `rval`
-            @is_thunk_type
-            def rval(
-                p=p, i=node_input_storage, o=node_output_storage, n=node, params=None
-            ):
-                r = p(n, [x[0] for x in i], o)
-                for o in node.outputs:
-                    compute_map[o][0] = True
-                return r
-
-        else:
-            params_val = node.params_type.filter(params)
-
-            @is_thunk_type
-            def rval(
-                p=p,
-                i=node_input_storage,
-                o=node_output_storage,
-                n=node,
-                params=params_val,
-            ):
-                r = p(n, [x[0] for x in i], o, params)
-                for o in node.outputs:
-                    compute_map[o][0] = True
-                return r
+        @is_thunk_type
+        def rval(p=p, i=node_input_storage, o=node_output_storage, n=node):
+            r = p(n, [x[0] for x in i], o)
+            for o in node.outputs:
+                compute_map[o][0] = True
+            return r
 
         rval.inputs = node_input_storage
         rval.outputs = node_output_storage
@@ -640,7 +593,7 @@ class _NoPythonOp(Op):
 
     """
 
-    def perform(self, node, inputs, output_storage, params=None):
+    def perform(self, node, inputs, output_storage):
         raise NotImplementedError("No Python implementation is provided by this Op.")
 
 

{pytensor-2.18.0 → 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.0 → pytensor-2.18.2}/pytensor/link/c/basic.py

@@ -20,6 +20,7 @@ from pytensor.graph.basic import (
     io_toposort,
     vars_between,
 )
+from pytensor.graph.utils import MethodNotDefined
 from pytensor.link.basic import Container, Linker, LocalLinker, PerformLinker
 from pytensor.link.c.cmodule import (
     METH_VARARGS,
@@ -617,7 +618,12 @@ class CLinker(Linker):
         # that needs it
         self.node_params = dict()
         for node in self.node_order:
-            params = node.run_params()
+            if not isinstance(node.op, CLinkerOp):
+                continue
+            try:
+                params = node.op.get_params(node)
+            except MethodNotDefined:
+                params = NoParams
             if params is not NoParams:
                 # try to avoid creating more than one variable for the
                 # same params.
@@ -803,7 +809,10 @@ class CLinker(Linker):
 
             sub = dict(failure_var=failure_var)
 
-            params = node.run_params()
+            try:
+                params = op.get_params(node)
+            except MethodNotDefined:
+                params = NoParams
             if params is not NoParams:
                 params_var = symbol[self.node_params[params]]
 

{pytensor-2.18.0 → 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
 
@@ -2744,7 +2745,9 @@ def default_blas_ldflags():
             [pathlib.Path(p).resolve() for p in line[len("libraries: =") :].split(":")]
             for line in stdout.decode(sys.stdout.encoding).splitlines()
             if line.startswith("libraries: =")
-        ][0]
+        ]
+        if len(maybe_lib_dirs) > 0:
+            maybe_lib_dirs = maybe_lib_dirs[0]
         return [str(d) for d in maybe_lib_dirs if d.exists() and d.is_dir()]
 
     def check_libs(
@@ -2777,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
@@ -2793,6 +2802,18 @@ def default_blas_ldflags():
 
     cxx_library_dirs = get_cxx_library_dirs()
     searched_library_dirs = cxx_library_dirs + _std_lib_dirs
+    if sys.platform == "win32":
+        # Conda on Windows saves MKL libraries under CONDA_PREFIX\Library\bin
+        # From the conda manual (https://docs.conda.io/projects/conda-build/en/stable/user-guide/environment-variables.html)
+        # it seems like conda could also save some libraries into the CONDA_PREFIX\Library\lib
+        # 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 [
@@ -2808,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=[
@@ -2820,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,
@@ -2840,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"],
@@ -2862,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.0 → pytensor-2.18.2}/pytensor/link/c/interface.py

@@ -1,11 +1,16 @@
+import typing
 import warnings
 from abc import abstractmethod
-from typing import Callable
+from typing import Callable, Optional
 
 from pytensor.graph.basic import Apply, Constant
 from pytensor.graph.utils import MethodNotDefined
 
 
+if typing.TYPE_CHECKING:
+    from pytensor.link.c.params_type import Params, ParamsType
+
+
 class CLinkerObject:
     """Standard methods for an `Op` or `Type` used with the `CLinker`."""
 
@@ -172,6 +177,8 @@ class CLinkerObject:
 class CLinkerOp(CLinkerObject):
     """Interface definition for `Op` subclasses compiled by `CLinker`."""
 
+    params_type: Optional["ParamsType"] = None
+
     @abstractmethod
     def c_code(
         self,
@@ -362,6 +369,22 @@ class CLinkerOp(CLinkerObject):
         """
         return ""
 
+    def get_params(self, node: Apply) -> "Params":
+        """Try to get parameters for the `Op` when :attr:`Op.params_type` is set to a `ParamsType`."""
+        if self.params_type is not None:
+            wrapper = self.params_type
+            if not all(hasattr(self, field) for field in wrapper.fields):
+                # Let's print missing attributes for debugging.
+                not_found = tuple(
+                    field for field in wrapper.fields if not hasattr(self, field)
+                )
+                raise AttributeError(
+                    f"{type(self).__name__}: missing attributes {not_found} for ParamsType."
+                )
+            # ParamsType.get_params() will apply filtering to attributes.
+            return self.params_type.get_params(self)
+        raise MethodNotDefined("get_params")
+
 
 class CLinkerType(CLinkerObject):
     r"""Interface specification for `Type`\s that can be arguments to a `CLinkerOp`.

{pytensor-2.18.0 → pytensor-2.18.2}/pytensor/link/c/op.py

@@ -664,7 +664,7 @@ class _NoPythonCOp(COp):
 
     """
 
-    def perform(self, node, inputs, output_storage, params=None):
+    def perform(self, node, inputs, output_storage):
         raise NotImplementedError("No Python implementation is provided by this COp.")
 
 
@@ -675,7 +675,7 @@ class _NoPythonExternalCOp(ExternalCOp):
 
     """
 
-    def perform(self, node, inputs, output_storage, params=None):
+    def perform(self, node, inputs, output_storage):
         raise NotImplementedError(
             "No Python implementation is provided by this ExternalCOp."
         )

{pytensor-2.18.0 → 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