pytensor 2.16.3__tar.gz → 2.17.1__tar.gz

{pytensor-2.16.3 → pytensor-2.17.1}/PKG-INFO

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

{pytensor-2.16.3 → pytensor-2.17.1}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2023-09-24T14:37:18+0200",
+ "date": "2023-10-04T13:04:02+0100",
  "dirty": false,
  "error": null,
- "full-revisionid": "2d944072a7cb37e931aa95d97f439f8b21004d4e",
- "version": "2.16.3"
+ "full-revisionid": "849c3b86553afd56081b0a2bd80701ea286acf02",
+ "version": "2.17.1"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.16.3 → pytensor-2.17.1}/pytensor/link/jax/dispatch/random.py

@@ -216,21 +216,20 @@ def jax_sample_fn_uniform(op):
 
 @jax_sample_fn.register(aer.ParetoRV)
 @jax_sample_fn.register(aer.GammaRV)
-def jax_sample_fn_shape_rate(op):
-    """JAX implementation of random variables in the shape-rate family.
+def jax_sample_fn_shape_scale(op):
+    """JAX implementation of random variables in the shape-scale family.
 
     JAX only implements the standard version of random variables in the
-    shape-rate family. We thus need to rescale the results manually.
+    shape-scale family. We thus need to rescale the results manually.
 
     """
     name = op.name
     jax_op = getattr(jax.random, name)
 
-    def sample_fn(rng, size, dtype, *parameters):
+    def sample_fn(rng, size, dtype, shape, scale):
         rng_key = rng["jax_state"]
         rng_key, sampling_key = jax.random.split(rng_key, 2)
-        (shape, rate) = parameters
-        sample = jax_op(sampling_key, shape, size, dtype) / rate
+        sample = jax_op(sampling_key, shape, size, dtype) * scale
         rng["jax_state"] = rng_key
         return (rng, sample)
 

{pytensor-2.16.3 → pytensor-2.17.1}/pytensor/tensor/nlinalg.py

@@ -541,9 +541,9 @@ class SVD(Op):
     # See doc in the docstring of the function just after this class.
     __props__ = ("full_matrices", "compute_uv")
 
-    def __init__(self, full_matrices=True, compute_uv=True):
-        self.full_matrices = full_matrices
-        self.compute_uv = compute_uv
+    def __init__(self, full_matrices: bool = True, compute_uv: bool = True):
+        self.full_matrices = bool(full_matrices)
+        self.compute_uv = bool(compute_uv)
 
     def make_node(self, x):
         x = as_tensor_variable(x)
@@ -584,7 +584,7 @@ class SVD(Op):
             return [s_shape]
 
 
-def svd(a, full_matrices=1, compute_uv=1):
+def svd(a, full_matrices: bool = True, compute_uv: bool = True):
     """
     This function performs the SVD on CPU.
 

{pytensor-2.16.3 → pytensor-2.17.1}/pytensor/tensor/random/basic.py

@@ -1,4 +1,5 @@
 import abc
+import warnings
 from typing import List, Optional, Union
 
 import numpy as np
@@ -419,7 +420,7 @@ class LogNormalRV(RandomVariable):
 lognormal = LogNormalRV()
 
 
-class GammaRV(ScipyRandomVariable):
+class GammaRV(RandomVariable):
     r"""A gamma continuous random variable.
 
     The probability density function for `gamma` in terms of the shape parameter
@@ -443,7 +444,7 @@ class GammaRV(ScipyRandomVariable):
     dtype = "floatX"
     _print_name = ("Gamma", "\\operatorname{Gamma}")
 
-    def __call__(self, shape, rate, size=None, **kwargs):
+    def __call__(self, shape, scale, size=None, **kwargs):
         r"""Draw samples from a gamma distribution.
 
         Signature
@@ -455,8 +456,8 @@ class GammaRV(ScipyRandomVariable):
         ----------
         shape
             The shape :math:`\alpha` of the gamma distribution. Must be positive.
-        rate
-            The rate :math:`\beta` of the gamma distribution. Must be positive.
+        scale
+            The scale :math:`1/\beta` of the gamma distribution. Must be positive.
         size
             Sample shape. If the given size is, e.g. `(m, n, k)` then `m * n * k`
             independent, identically distributed random variables are
@@ -464,14 +465,26 @@ class GammaRV(ScipyRandomVariable):
             is returned.
 
         """
-        return super().__call__(shape, 1.0 / rate, size=size, **kwargs)
+        return super().__call__(shape, scale, size=size, **kwargs)
 
-    @classmethod
-    def rng_fn_scipy(cls, rng, shape, scale, size):
-        return stats.gamma.rvs(shape, scale=scale, size=size, random_state=rng)
 
+_gamma = GammaRV()
+
+
+def gamma(shape, rate=None, scale=None, **kwargs):
+    # TODO: Remove helper when rate is deprecated
+    if rate is not None and scale is not None:
+        raise ValueError("Cannot specify both rate and scale")
+    elif rate is None and scale is None:
+        raise ValueError("Must specify scale")
+    elif rate is not None:
+        warnings.warn(
+            "Gamma rate argument is deprecated and will stop working, use scale instead",
+            FutureWarning,
+        )
+        scale = 1.0 / rate
 
-gamma = GammaRV()
+    return _gamma(shape, scale, **kwargs)
 
 
 class ChiSquareRV(RandomVariable):

{pytensor-2.16.3 → pytensor-2.17.1}/pytensor/tensor/random/rewriting/jax.py

@@ -15,9 +15,9 @@ from pytensor.tensor.random.basic import (
     LogNormalRV,
     NegBinomialRV,
     WaldRV,
+    _gamma,
     beta,
     binomial,
-    gamma,
     normal,
     poisson,
     uniform,
@@ -92,7 +92,7 @@ def geometric_from_uniform(fgraph, node):
 @node_rewriter([NegBinomialRV])
 def negative_binomial_from_gamma_poisson(fgraph, node):
     rng, *other_inputs, n, p = node.inputs
-    next_rng, g = gamma.make_node(rng, *other_inputs, n, p / (1 - p)).outputs
+    next_rng, g = _gamma.make_node(rng, *other_inputs, n, (1 - p) / p).outputs
     next_rng, p = poisson.make_node(next_rng, *other_inputs, g).outputs
     return [next_rng, p]
 
@@ -100,21 +100,21 @@ def negative_binomial_from_gamma_poisson(fgraph, node):
 @node_rewriter([InvGammaRV])
 def inverse_gamma_from_gamma(fgraph, node):
     *other_inputs, shape, scale = node.inputs
-    next_rng, g = gamma.make_node(*other_inputs, shape, scale).outputs
+    next_rng, g = _gamma.make_node(*other_inputs, shape, 1 / scale).outputs
     return [next_rng, reciprocal(g)]
 
 
 @node_rewriter([ChiSquareRV])
 def chi_square_from_gamma(fgraph, node):
     *other_inputs, df = node.inputs
-    next_rng, g = gamma.make_node(*other_inputs, df / 2, 1 / 2).outputs
+    next_rng, g = _gamma.make_node(*other_inputs, df / 2, 2).outputs
     return [next_rng, g]
 
 
 @node_rewriter([GenGammaRV])
 def generalized_gamma_from_gamma(fgraph, node):
     *other_inputs, alpha, p, lambd = node.inputs
-    next_rng, g = gamma.make_node(*other_inputs, alpha / p, ones_like(lambd)).outputs
+    next_rng, g = _gamma.make_node(*other_inputs, alpha / p, ones_like(lambd)).outputs
     g = (g ** reciprocal(p)) * lambd
     return [next_rng, cast(g, dtype=node.default_output().dtype)]
 

{pytensor-2.16.3 → pytensor-2.17.1}/pytensor/tensor/rewriting/linalg.py

@@ -6,7 +6,7 @@ from pytensor.tensor.basic import TensorVariable, diagonal, swapaxes
 from pytensor.tensor.blas import Dot22
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.elemwise import DimShuffle
-from pytensor.tensor.math import Dot, Prod, log, prod
+from pytensor.tensor.math import Dot, Prod, _matrix_matrix_matmul, log, prod
 from pytensor.tensor.nlinalg import MatrixInverse, det
 from pytensor.tensor.rewriting.basic import (
     register_canonicalize,
@@ -168,13 +168,25 @@ def cholesky_ldotlt(fgraph, node):
     rewrite cholesky(dot(L, L.T), lower=True) = L, where L is lower triangular,
     or cholesky(dot(U.T, U), upper=True) = U where U is upper triangular.
 
+    Also works with matmul.
+
     This utilizes a boolean `lower_triangular` or `upper_triangular` tag on matrices.
     """
     if not isinstance(node.op.core_op, Cholesky):
         return
 
     A = node.inputs[0]
-    if not (A.owner and isinstance(A.owner.op, (Dot, Dot22))):
+    if not (
+        A.owner is not None
+        and (
+            (
+                isinstance(A.owner.op, (Dot, Dot22))
+                # This rewrite only applies to matrix Dot
+                and A.owner.inputs[0].type.ndim == 2
+            )
+            or (A.owner.op == _matrix_matrix_matmul)
+        )
+    ):
         return
 
     l, r = A.owner.inputs

{pytensor-2.16.3 → pytensor-2.17.1}/pytensor/tensor/rewriting/math.py

@@ -2081,63 +2081,65 @@ def local_pow_to_nested_squaring(fgraph, node):
     Note: This sounds like the kind of thing any half-decent compiler can do by itself?
     """
 
-    if node.op == at_pow:
-        # the idea here is that we have pow(x, y)
-        odtype = node.outputs[0].dtype
-        xsym = node.inputs[0]
-        ysym = node.inputs[1]
-        y = get_constant(ysym)
-
-        # the next line is needed to fix a strange case that I don't
-        # know how to make a separate test.
-        # That happen in the `test_log_erfc` test.
-        # y is a ndarray with dtype int8 and value 2,4 or 6. This make
-        # the abs(y) <= 512 fail!
-        # taking the value outside ndarray solve the problem.
-        # it could be that in that case, numpy make the comparison
-        # into the wrong type(do in int8 that overflow.)
-        if isinstance(y, np.ndarray):
-            assert y.size == 1
-            try:
-                y = y[0]
-            except IndexError:
-                pass
-        if (y is not None) and not broadcasted_by(xsym, ysym):
-            rval = None
-            # 512 is too small for the cpu and too big for some gpu!
-            if abs(y) == int(abs(y)) and abs(y) <= 512:
-                pow2 = [xsym]
-                pow2_scal = [aes.get_scalar_type(xsym.dtype)()]
-                y_to_do = abs(y)
-                for i in range(int(np.log2(y_to_do))):
-                    pow2.append(sqr(pow2[i]))
-                    pow2_scal.append(aes.sqr(pow2_scal[i]))
-                rval1 = None
-                rval1_scal = None
-                while y_to_do > 0:
-                    log_to_do = int(np.log2(y_to_do))
-                    if rval1:
-                        rval1 *= pow2[log_to_do]
-                        rval1_scal *= pow2_scal[log_to_do]
-                    else:
-                        rval1 = pow2[log_to_do]
-                        rval1_scal = pow2_scal[log_to_do]
-                    y_to_do -= 2**log_to_do
-
-                if abs(y) > 2:
-                    # We fuse all the pow together here to make
-                    # compilation faster
-                    rval1 = Elemwise(
-                        aes.Composite([pow2_scal[0]], [rval1_scal])
-                    ).make_node(xsym)
-                if y < 0:
-                    rval = [reciprocal(rval1)]
+    # the idea here is that we have pow(x, y)
+    odtype = node.outputs[0].dtype
+    xsym = node.inputs[0]
+    ysym = node.inputs[1]
+    y = get_constant(ysym)
+
+    # the next line is needed to fix a strange case that I don't
+    # know how to make a separate test.
+    # That happen in the `test_log_erfc` test.
+    # y is a ndarray with dtype int8 and value 2,4 or 6. This make
+    # the abs(y) <= 512 fail!
+    # taking the value outside ndarray solve the problem.
+    # it could be that in that case, numpy make the comparison
+    # into the wrong type(do in int8 that overflow.)
+    if isinstance(y, np.ndarray):
+        assert y.size == 1
+        try:
+            y = y[0]
+        except IndexError:
+            pass
+    if (y is not None) and not broadcasted_by(xsym, ysym):
+        rval = None
+        # 512 is too small for the cpu and too big for some gpu!
+        if abs(y) == int(abs(y)) and abs(y) <= 512:
+            pow2 = [xsym]
+            pow2_scal = [aes.get_scalar_type(xsym.dtype)()]
+            y_to_do = abs(y)
+            for i in range(int(np.log2(y_to_do))):
+                pow2.append(sqr(pow2[i]))
+                pow2_scal.append(aes.sqr(pow2_scal[i]))
+            rval1 = None
+            rval1_scal = None
+            while y_to_do > 0:
+                log_to_do = int(np.log2(y_to_do))
+                if rval1:
+                    rval1 *= pow2[log_to_do]
+                    rval1_scal *= pow2_scal[log_to_do]
                 else:
-                    rval = [rval1]
-            if rval:
-                rval[0] = cast(rval[0], odtype)
-                assert rval[0].type == node.outputs[0].type, (rval, node.outputs)
-                return rval
+                    rval1 = pow2[log_to_do]
+                    rval1_scal = pow2_scal[log_to_do]
+                y_to_do -= 2**log_to_do
+
+            if abs(y) > 2:
+                # We fuse all the pow together here to make
+                # compilation faster
+                rval1 = Elemwise(aes.Composite([pow2_scal[0]], [rval1_scal])).make_node(
+                    xsym
+                )
+            if y < 0:
+                rval = [reciprocal(rval1)]
+            else:
+                rval = [rval1]
+        if rval:
+            rval[0] = cast(rval[0], odtype)
+            # TODO: We can add a specify_broadcastable and/or unbroadcast to make the
+            #  output types compatible. Or work on #408 and let TensorType.filter_variable do it.
+            if rval[0].type.broadcastable != node.outputs[0].type.broadcastable:
+                return None
+            return rval
 
 
 @register_specialize

{pytensor-2.16.3 → pytensor-2.17.1}/pytensor/tensor/variable.py

@@ -652,7 +652,7 @@ class _tensor_py_operators:
         return at.math.matmul(left, right)
 
     def __rmatmul__(right, left):
-        return at.math.matmul(right, left)
+        return at.math.matmul(left, right)
 
     def sum(self, axis=None, dtype=None, keepdims=False, acc_dtype=None):
         """See :func:`pytensor.tensor.math.sum`."""

{pytensor-2.16.3 → pytensor-2.17.1}/pytensor.egg-info/PKG-INFO

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