Merge branch 'main' of github.com:pymc-devs/pytensor into elemwise_torch_improvement

Author: Ch0ronomato

.pre-commit-config.yaml

@@ -9,7 +9,7 @@ exclude: |
     )$
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.6.0
+    rev: v5.0.0
     hooks:
       - id: debug-statements
         exclude: |
@@ -27,7 +27,7 @@ repos:
     - id: sphinx-lint 
       args: ["."]
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.6.5
+    rev: v0.7.1
     hooks:
       - id: ruff
         args: ["--fix", "--output-format=full"]

pytensor/compile/function/types.py

@@ -128,9 +128,6 @@ def fiter_variable(self, other):
             " a symbolic placeholder."
         )
 
-    def may_share_memory(a, b):
-        return False
-
     def value_eq(a, b, force_same_dtype=True):
         raise AssertionError(
             "If you're assigning to a DisconnectedType you're"

pytensor/gradient.py

@@ -26,9 +26,6 @@ def filter(self, data, strict=False, allow_downcast=None):
     def filter_variable(self, other, allow_convert=True):
         raise ValueError("No values may be assigned to a NullType")
 
-    def may_share_memory(a, b):
-        return False
-
     def values_eq(self, a, b, force_same_dtype=True):
         raise ValueError("NullType has no values to compare")
 

pytensor/graph/null_type.py

@@ -513,17 +513,24 @@ def make_py_thunk(
         """
         node_input_storage = [storage_map[r] for r in node.inputs]
         node_output_storage = [storage_map[r] for r in node.outputs]
+        node_compute_map = [compute_map[r] for r in node.outputs]
 
         if debug and hasattr(self, "debug_perform"):
             p = node.op.debug_perform
         else:
             p = node.op.perform
 
         @is_thunk_type
-        def rval(p=p, i=node_input_storage, o=node_output_storage, n=node):
+        def rval(
+            p=p,
+            i=node_input_storage,
+            o=node_output_storage,
+            n=node,
+            cm=node_compute_map,
+        ):
             r = p(n, [x[0] for x in i], o)
-            for o in node.outputs:
-                compute_map[o][0] = True
+            for entry in cm:
+                entry[0] = True
             return r
 
         rval.inputs = node_input_storage

pytensor/graph/op.py

@@ -48,10 +48,7 @@ def in_same_class(self, otype: "Type") -> bool | None:
         unique element (i.e. it uses `self.__eq__`).
 
         """
-        if self == otype:
-            return True
-
-        return False
+        return self == otype
 
     def is_super(self, otype: "Type") -> bool | None:
         """Determine if `self` is a supertype of `otype`.

pytensor/graph/type.py

@@ -2007,13 +2007,18 @@ def try_blas_flag(flags):
     cflags.extend(f"-L{path_wrapper}{d}{path_wrapper}" for d in std_lib_dirs())
 
     res = GCC_compiler.try_compile_tmp(
-        test_code, tmp_prefix="try_blas_", flags=cflags, try_run=True
+        test_code, tmp_prefix="try_blas_", flags=cflags, try_run=True, output=True
     )
     # res[0]: shows successful compilation
     # res[1]: shows successful execution
+    # res[2]: shows execution results
+    # res[3]: shows execution or compilation error message
     if res and res[0] and res[1]:
         return " ".join(flags)
     else:
+        _logger.debug(
+            "try_blas_flags of flags: %r\nfailed with error message %s", flags, res[3]
+        )
         return ""
 
 
@@ -2801,7 +2806,6 @@ def check_libs(
             _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")
 

pytensor/link/c/cmodule.py

@@ -303,13 +303,6 @@ def clone(self, dtype=None, **kwargs):
             dtype = self.dtype
         return type(self)(dtype)
 
-    @staticmethod
-    def may_share_memory(a, b):
-        # This class represent basic c type, represented in python
-        # with numpy.scalar. They are read only. So from python, they
-        # can never share memory.
-        return False
-
     def filter(self, data, strict=False, allow_downcast=None):
         py_type = self.dtype_specs()[0]
         if strict and not isinstance(data, py_type):
@@ -4253,7 +4246,11 @@ def __str__(self):
             r.name = f"o{int(i)}"
         io = set(self.fgraph.inputs + self.fgraph.outputs)
         for i, r in enumerate(self.fgraph.variables):
-            if r not in io and len(self.fgraph.clients[r]) > 1:
+            if (
+                not isinstance(r, Constant)
+                and r not in io
+                and len(self.fgraph.clients[r]) > 1
+            ):
                 r.name = f"t{int(i)}"
 
         if len(self.fgraph.outputs) > 1 or len(self.fgraph.apply_nodes) > 10:
@@ -4352,7 +4349,7 @@ def c_code_template(self):
                 if var not in self.fgraph.inputs:
                     # This is an orphan
                     if isinstance(var, Constant) and isinstance(var.type, CLinkerType):
-                        subd[var] = var.type.c_literal(var.data)
+                        subd[var] = f"({var.type.c_literal(var.data)})"
                     else:
                         raise ValueError(
                             "All orphans in the fgraph to Composite must"
@@ -4411,7 +4408,7 @@ def c_code(self, node, nodename, inames, onames, sub):
         return self.c_code_template % d
 
     def c_code_cache_version_outer(self) -> tuple[int, ...]:
-        return (4,)
+        return (5,)
 
 
 class Compositef32:

pytensor/scalar/basic.py

@@ -239,7 +239,7 @@ def c_code_template(self):
                 if var not in self.fgraph.inputs:
                     # This is an orphan
                     if isinstance(var, Constant) and isinstance(var.type, CLinkerType):
-                        subd[var] = var.type.c_literal(var.data)
+                        subd[var] = f"({var.type.c_literal(var.data)})"
                     else:
                         raise ValueError(
                             "All orphans in the fgraph to ScalarLoop must"
@@ -342,4 +342,4 @@ def c_code(self, node, nodename, inames, onames, sub):
         return res
 
     def c_code_cache_version_outer(self):
-        return (2,)
+        return (3,)

pytensor/scalar/loop.py

@@ -1,5 +1,4 @@
 from collections.abc import Sequence
-from copy import copy
 from typing import Any, cast
 
 import numpy as np
@@ -79,7 +78,6 @@ def __init__(
         self.name = name
         self.inputs_sig, self.outputs_sig = _parse_gufunc_signature(signature)
         self.gufunc_spec = gufunc_spec
-        self._gufunc = None
         if destroy_map is not None:
             self.destroy_map = destroy_map
         if self.destroy_map != core_op.destroy_map:
@@ -91,11 +89,6 @@ def __init__(
 
         super().__init__(**kwargs)
 
-    def __getstate__(self):
-        d = copy(self.__dict__)
-        d["_gufunc"] = None
-        return d
-
     def _create_dummy_core_node(self, inputs: Sequence[TensorVariable]) -> Apply:
         core_input_types = []
         for i, (inp, sig) in enumerate(zip(inputs, self.inputs_sig)):
@@ -296,32 +289,46 @@ def L_op(self, inputs, outs, ograds):
 
         return rval
 
-    def _create_gufunc(self, node):
+    def _create_node_gufunc(self, node) -> None:
+        """Define (or retrieve) the node gufunc used in `perform`.
+
+        If the Blockwise or core_op have a `gufunc_spec`, the relevant numpy or scipy gufunc is used directly.
+        Otherwise, we default to `np.vectorize` of the core_op `perform` method for a dummy node.
+
+        The gufunc is stored in the tag of the node.
+        """
         gufunc_spec = self.gufunc_spec or getattr(self.core_op, "gufunc_spec", None)
 
         if gufunc_spec is not None:
-            self._gufunc = import_func_from_string(gufunc_spec[0])
-            if self._gufunc:
-                return self._gufunc
-            else:
+            gufunc = import_func_from_string(gufunc_spec[0])
+            if gufunc is None:
                 raise ValueError(f"Could not import gufunc {gufunc_spec[0]} for {self}")
 
-        n_outs = len(self.outputs_sig)
-        core_node = self._create_dummy_core_node(node.inputs)
-
-        def core_func(*inner_inputs):
-            inner_outputs = [[None] for _ in range(n_outs)]
+        else:
+            # Wrap core_op perform method in numpy vectorize
+            n_outs = len(self.outputs_sig)
+            core_node = self._create_dummy_core_node(node.inputs)
+            inner_outputs_storage = [[None] for _ in range(n_outs)]
+
+            def core_func(
+                *inner_inputs,
+                core_node=core_node,
+                inner_outputs_storage=inner_outputs_storage,
+            ):
+                self.core_op.perform(
+                    core_node,
+                    [np.asarray(inp) for inp in inner_inputs],
+                    inner_outputs_storage,
+                )
 
-            inner_inputs = [np.asarray(inp) for inp in inner_inputs]
-            self.core_op.perform(core_node, inner_inputs, inner_outputs)
+                if n_outs == 1:
+                    return inner_outputs_storage[0][0]
+                else:
+                    return tuple(r[0] for r in inner_outputs_storage)
 
-            if len(inner_outputs) == 1:
-                return inner_outputs[0][0]
-            else:
-                return tuple(r[0] for r in inner_outputs)
+            gufunc = np.vectorize(core_func, signature=self.signature)
 
-        self._gufunc = np.vectorize(core_func, signature=self.signature)
-        return self._gufunc
+        node.tag.gufunc = gufunc
 
     def _check_runtime_broadcast(self, node, inputs):
         batch_ndim = self.batch_ndim(node)
@@ -340,10 +347,12 @@ def _check_runtime_broadcast(self, node, inputs):
                 )
 
     def perform(self, node, inputs, output_storage):
-        gufunc = self._gufunc
+        gufunc = getattr(node.tag, "gufunc", None)
 
         if gufunc is None:
-            gufunc = self._create_gufunc(node)
+            # Cache it once per node
+            self._create_node_gufunc(node)
+            gufunc = node.tag.gufunc
 
         self._check_runtime_broadcast(node, inputs)
 

pytensor/tensor/blockwise.py

@@ -387,24 +387,17 @@ def dist_params(self, node) -> Sequence[Variable]:
         return node.inputs[2:]
 
     def perform(self, node, inputs, outputs):
-        rng_var_out, smpl_out = outputs
-
         rng, size, *args = inputs
 
         # Draw from `rng` if `self.inplace` is `True`, and from a copy of `rng` otherwise.
         if not self.inplace:
             rng = copy(rng)
 
-        rng_var_out[0] = rng
-
-        if size is not None:
-            size = tuple(size)
-        smpl_val = self.rng_fn(rng, *([*args, size]))
-
-        if not isinstance(smpl_val, np.ndarray) or str(smpl_val.dtype) != self.dtype:
-            smpl_val = np.asarray(smpl_val, dtype=self.dtype)
-
-        smpl_out[0] = smpl_val
+        outputs[0][0] = rng
+        outputs[1][0] = np.asarray(
+            self.rng_fn(rng, *args, None if size is None else tuple(size)),
+            dtype=self.dtype,
+        )
 
     def grad(self, inputs, outputs):
         return [

pytensor/tensor/random/op.py

@@ -127,8 +127,8 @@ def local_blockwise_alloc(fgraph, node):
                 value, *shape = inp.owner.inputs
 
                 # Check what to do with the value of the Alloc
-                squeezed_value = _squeeze_left(value, batch_ndim)
-                missing_ndim = len(shape) - value.type.ndim
+                missing_ndim = inp.type.ndim - value.type.ndim
+                squeezed_value = _squeeze_left(value, (batch_ndim - missing_ndim))
                 if (
                     (((1,) * missing_ndim + value.type.broadcastable)[batch_ndim:])
                     != inp.type.broadcastable[batch_ndim:]

pytensor/tensor/rewriting/blockwise.py

@@ -4,9 +4,11 @@
 
 from pytensor import Variable
 from pytensor import tensor as pt
+from pytensor.compile import optdb
 from pytensor.graph import Apply, FunctionGraph
 from pytensor.graph.rewriting.basic import (
     copy_stack_trace,
+    in2out,
     node_rewriter,
 )
 from pytensor.scalar.basic import Mul
@@ -45,9 +47,11 @@
     Cholesky,
     Solve,
     SolveBase,
+    _bilinear_solve_discrete_lyapunov,
     block_diag,
     cholesky,
     solve,
+    solve_discrete_lyapunov,
     solve_triangular,
 )
 
@@ -966,3 +970,22 @@ def rewrite_cholesky_diag_to_sqrt_diag(fgraph, node):
             non_eye_input = pt.shape_padaxis(non_eye_input, -2)
 
     return [eye_input * (non_eye_input**0.5)]
+
+
+@node_rewriter([_bilinear_solve_discrete_lyapunov])
+def jax_bilinaer_lyapunov_to_direct(fgraph: FunctionGraph, node: Apply):
+    """
+    Replace BilinearSolveDiscreteLyapunov with a direct computation that is supported by JAX
+    """
+    A, B = (cast(TensorVariable, x) for x in node.inputs)
+    result = solve_discrete_lyapunov(A, B, method="direct")
+
+    return [result]
+
+
+optdb.register(
+    "jax_bilinaer_lyapunov_to_direct",
+    in2out(jax_bilinaer_lyapunov_to_direct),
+    "jax",
+    position=0.9,  # Run before canonicalization
+)