pytensor 2.21.0__tar.gz → 2.22.1__tar.gz

{pytensor-2.21.0/pytensor.egg-info → pytensor-2.22.1}/PKG-INFO

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

{pytensor-2.21.0 → pytensor-2.22.1}/pytensor/__init__.py

@@ -159,6 +159,7 @@ def get_underlying_scalar_constant(v):
 # isort: off
 import pytensor.tensor.random.var
 import pytensor.sparse
+from pytensor.ifelse import ifelse
 from pytensor.scan import checkpoints
 from pytensor.scan.basic import scan
 from pytensor.scan.views import foldl, foldr, map, reduce

{pytensor-2.21.0 → pytensor-2.22.1}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-05-13T21:55:31+0200",
+ "date": "2024-05-28T12:10:34+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "90da9e61688b0fef8b91106f0e2ae08ee67be1be",
- "version": "2.21.0"
+ "full-revisionid": "fc21336a67334241d1c9a59383084ce453988d4e",
+ "version": "2.22.1"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.21.0 → pytensor-2.22.1}/pytensor/graph/basic.py

@@ -266,7 +266,7 @@ class Apply(Node, Generic[OpType]):
 
         assert isinstance(inputs, list | tuple)
         remake_node = False
-        new_inputs: list["Variable"] = list(inputs)
+        new_inputs: list[Variable] = list(inputs)
 
         # Some Ops like Alloc require the node to always be rebuilt in non-strict mode
         # as the output type depends on the input values and not just their types
@@ -1272,7 +1272,7 @@ def clone_node_and_cache(
         return None
 
     # Use a cached `Op` clone when available
-    new_op: "Op" | None = cast(Optional["Op"], clone_d.get(node.op))
+    new_op: Op | None = cast(Optional["Op"], clone_d.get(node.op))
 
     cloned_inputs: list[Variable] = [cast(Variable, clone_d[i]) for i in node.inputs]
 
@@ -1789,38 +1789,6 @@ def view_roots(node: Variable) -> list[Variable]:
         return [node]
 
 
-def list_of_nodes(
-    inputs: Collection[Variable], outputs: Iterable[Variable]
-) -> list[Apply]:
-    r"""Return the `Apply` nodes of the graph between `inputs` and `outputs`.
-
-    Parameters
-    ----------
-    inputs : list of Variable
-        Input `Variable`\s.
-    outputs : list of Variable
-        Output `Variable`\s.
-
-    """
-
-    def expand(o: Apply) -> list[Apply]:
-        return [
-            inp.owner
-            for inp in o.inputs
-            if inp.owner and not any(i in inp.owner.outputs for i in inputs)
-        ]
-
-    return list(
-        cast(
-            Iterable[Apply],
-            walk(
-                [o.owner for o in outputs if o.owner],
-                expand,
-            ),
-        )
-    )
-
-
 def apply_depends_on(apply: Apply, depends_on: Apply | Collection[Apply]) -> bool:
     """Determine if any `depends_on` is in the graph given by ``apply``.
 

{pytensor-2.21.0 → pytensor-2.22.1}/pytensor/graph/rewriting/basic.py

@@ -1577,7 +1577,7 @@ class PatternNodeRewriter(NodeRewriter):
         """
         from pytensor.graph.rewriting.unify import convert_strs_to_vars
 
-        var_map: dict[str, "Var"] = {}
+        var_map: dict[str, Var] = {}
         self.in_pattern = convert_strs_to_vars(in_pattern, var_map=var_map)
         self.out_pattern = convert_strs_to_vars(out_pattern, var_map=var_map)
         self.values_eq_approx = values_eq_approx

{pytensor-2.21.0 → pytensor-2.22.1}/pytensor/link/jax/dispatch/random.py

@@ -8,6 +8,7 @@ from numpy.random.bit_generator import (  # type: ignore[attr-defined]
 )
 
 import pytensor.tensor.random.basic as ptr
+from pytensor.graph import Constant
 from pytensor.link.jax.dispatch.basic import jax_funcify, jax_typify
 from pytensor.link.jax.dispatch.shape import JAXShapeTuple
 from pytensor.tensor.shape import Shape, Shape_i
@@ -91,33 +92,46 @@ def jax_funcify_RandomVariable(op, node, **kwargs):
     """JAX implementation of random variables."""
     rv = node.outputs[1]
     out_dtype = rv.type.dtype
-    out_size = rv.type.shape
-
-    if op.ndim_supp > 0:
-        out_size = node.outputs[1].type.shape[: -op.ndim_supp]
+    static_shape = rv.type.shape
+
+    batch_ndim = op.batch_ndim(node)
+
+    # Try to pass static size directly to JAX
+    static_size = static_shape[:batch_ndim]
+    if None in static_size:
+        # Sometimes size can be constant folded during rewrites,
+        # without the RandomVariable node being updated with new static types
+        size_param = node.inputs[1]
+        if isinstance(size_param, Constant):
+            size_tuple = tuple(size_param.data)
+            # PyTensor uses empty size to represent size = None
+            if len(size_tuple):
+                static_size = tuple(size_param.data)
 
     # If one dimension has unknown size, either the size is determined
     # by a `Shape` operator in which case JAX will compile, or it is
     # not and we fail gracefully.
-    if None in out_size:
+    if None in static_size:
         assert_size_argument_jax_compatible(node)
 
         def sample_fn(rng, size, dtype, *parameters):
             # PyTensor uses empty size to represent size = None
             if jax.numpy.asarray(size).shape == (0,):
                 size = None
-            return jax_sample_fn(op)(rng, size, out_dtype, *parameters)
+            return jax_sample_fn(op, node=node)(rng, size, out_dtype, *parameters)
 
     else:
 
         def sample_fn(rng, size, dtype, *parameters):
-            return jax_sample_fn(op)(rng, out_size, out_dtype, *parameters)
+            return jax_sample_fn(op, node=node)(
+                rng, static_size, out_dtype, *parameters
+            )
 
     return sample_fn
 
 
 @singledispatch
-def jax_sample_fn(op):
+def jax_sample_fn(op, node):
     name = op.name
     raise NotImplementedError(
         f"No JAX implementation for the given distribution: {name}"
@@ -128,7 +142,7 @@ def jax_sample_fn(op):
 @jax_sample_fn.register(ptr.DirichletRV)
 @jax_sample_fn.register(ptr.PoissonRV)
 @jax_sample_fn.register(ptr.MvNormalRV)
-def jax_sample_fn_generic(op):
+def jax_sample_fn_generic(op, node):
     """Generic JAX implementation of random variables."""
     name = op.name
     jax_op = getattr(jax.random, name)
@@ -149,7 +163,7 @@ def jax_sample_fn_generic(op):
 @jax_sample_fn.register(ptr.LogisticRV)
 @jax_sample_fn.register(ptr.NormalRV)
 @jax_sample_fn.register(ptr.StandardNormalRV)
-def jax_sample_fn_loc_scale(op):
+def jax_sample_fn_loc_scale(op, node):
     """JAX implementation of random variables in the loc-scale families.
 
     JAX only implements the standard version of random variables in the
@@ -174,7 +188,7 @@ def jax_sample_fn_loc_scale(op):
 
 
 @jax_sample_fn.register(ptr.BernoulliRV)
-def jax_sample_fn_bernoulli(op):
+def jax_sample_fn_bernoulli(op, node):
     """JAX implementation of `BernoulliRV`."""
 
     # We need a separate dispatch, because there is no dtype argument for Bernoulli in JAX
@@ -189,7 +203,7 @@ def jax_sample_fn_bernoulli(op):
 
 
 @jax_sample_fn.register(ptr.CategoricalRV)
-def jax_sample_fn_categorical(op):
+def jax_sample_fn_categorical(op, node):
     """JAX implementation of `CategoricalRV`."""
 
     # We need a separate dispatch because Categorical expects logits in JAX
@@ -208,7 +222,7 @@ def jax_sample_fn_categorical(op):
 @jax_sample_fn.register(ptr.RandIntRV)
 @jax_sample_fn.register(ptr.IntegersRV)
 @jax_sample_fn.register(ptr.UniformRV)
-def jax_sample_fn_uniform(op):
+def jax_sample_fn_uniform(op, node):
     """JAX implementation of random variables with uniform density.
 
     We need to pass the arguments as keyword arguments since the order
@@ -236,7 +250,7 @@ def jax_sample_fn_uniform(op):
 
 @jax_sample_fn.register(ptr.ParetoRV)
 @jax_sample_fn.register(ptr.GammaRV)
-def jax_sample_fn_shape_scale(op):
+def jax_sample_fn_shape_scale(op, node):
     """JAX implementation of random variables in the shape-scale family.
 
     JAX only implements the standard version of random variables in the
@@ -259,7 +273,7 @@ def jax_sample_fn_shape_scale(op):
 
 
 @jax_sample_fn.register(ptr.ExponentialRV)
-def jax_sample_fn_exponential(op):
+def jax_sample_fn_exponential(op, node):
     """JAX implementation of `ExponentialRV`."""
 
     def sample_fn(rng, size, dtype, scale):
@@ -275,7 +289,7 @@ def jax_sample_fn_exponential(op):
 
 
 @jax_sample_fn.register(ptr.StudentTRV)
-def jax_sample_fn_t(op):
+def jax_sample_fn_t(op, node):
     """JAX implementation of `StudentTRV`."""
 
     def sample_fn(rng, size, dtype, df, loc, scale):
@@ -290,30 +304,111 @@ def jax_sample_fn_t(op):
     return sample_fn
 
 
-@jax_sample_fn.register(ptr.ChoiceRV)
-def jax_funcify_choice(op):
+@jax_sample_fn.register(ptr.ChoiceWithoutReplacement)
+def jax_funcify_choice(op, node):
     """JAX implementation of `ChoiceRV`."""
 
+    batch_ndim = op.batch_ndim(node)
+    a, *p, core_shape = node.inputs[3:]
+    a_core_ndim, *p_core_ndim, _ = op.ndims_params
+
+    if batch_ndim and a_core_ndim == 0:
+        raise NotImplementedError(
+            "Batch dimensions are not supported for 0d arrays. "
+            "A default JAX rewrite should have materialized the implicit arange"
+        )
+
+    a_batch_ndim = a.type.ndim - a_core_ndim
+    if op.has_p_param:
+        [p] = p
+        [p_core_ndim] = p_core_ndim
+        p_batch_ndim = p.type.ndim - p_core_ndim
+
     def sample_fn(rng, size, dtype, *parameters):
         rng_key = rng["jax_state"]
         rng_key, sampling_key = jax.random.split(rng_key, 2)
-        (a, p, replace) = parameters
-        smpl_value = jax.random.choice(sampling_key, a, size, replace, p)
+
+        if op.has_p_param:
+            a, p, core_shape = parameters
+        else:
+            a, core_shape = parameters
+            p = None
+        core_shape = tuple(np.asarray(core_shape))
+
+        if batch_ndim == 0:
+            sample = jax.random.choice(
+                sampling_key, a, shape=core_shape, replace=False, p=p
+            )
+
+        else:
+            if size is None:
+                if p is None:
+                    size = a.shape[:a_batch_ndim]
+                else:
+                    size = jax.numpy.broadcast_shapes(
+                        a.shape[:a_batch_ndim],
+                        p.shape[:p_batch_ndim],
+                    )
+
+            a = jax.numpy.broadcast_to(a, size + a.shape[a_batch_ndim:])
+            if p is not None:
+                p = jax.numpy.broadcast_to(p, size + p.shape[p_batch_ndim:])
+
+            batch_sampling_keys = jax.random.split(sampling_key, np.prod(size))
+
+            # Ravel the batch dimensions because vmap only works along a single axis
+            raveled_batch_a = a.reshape((-1,) + a.shape[batch_ndim:])
+            if p is None:
+                raveled_sample = jax.vmap(
+                    lambda key, a: jax.random.choice(
+                        key, a, shape=core_shape, replace=False, p=None
+                    )
+                )(batch_sampling_keys, raveled_batch_a)
+            else:
+                raveled_batch_p = p.reshape((-1,) + p.shape[batch_ndim:])
+                raveled_sample = jax.vmap(
+                    lambda key, a, p: jax.random.choice(
+                        key, a, shape=core_shape, replace=False, p=p
+                    )
+                )(batch_sampling_keys, raveled_batch_a, raveled_batch_p)
+
+            # Reshape the batch dimensions
+            sample = raveled_sample.reshape(size + raveled_sample.shape[1:])
+
         rng["jax_state"] = rng_key
-        return (rng, smpl_value)
+        return (rng, sample)
 
     return sample_fn
 
 
 @jax_sample_fn.register(ptr.PermutationRV)
-def jax_sample_fn_permutation(op):
+def jax_sample_fn_permutation(op, node):
     """JAX implementation of `PermutationRV`."""
 
+    batch_ndim = op.batch_ndim(node)
+    x_batch_ndim = node.inputs[-1].type.ndim - op.ndims_params[0]
+
     def sample_fn(rng, size, dtype, *parameters):
         rng_key = rng["jax_state"]
         rng_key, sampling_key = jax.random.split(rng_key, 2)
         (x,) = parameters
-        sample = jax.random.permutation(sampling_key, x)
+        if batch_ndim:
+            # jax.random.permutation has no concept of batch dims
+            x_core_shape = x.shape[x_batch_ndim:]
+            if size is None:
+                size = x.shape[:x_batch_ndim]
+            else:
+                x = jax.numpy.broadcast_to(x, size + x_core_shape)
+
+            batch_sampling_keys = jax.random.split(sampling_key, np.prod(size))
+            raveled_batch_x = x.reshape((-1,) + x.shape[batch_ndim:])
+            raveled_sample = jax.vmap(lambda key, x: jax.random.permutation(key, x))(
+                batch_sampling_keys, raveled_batch_x
+            )
+            sample = raveled_sample.reshape(size + raveled_sample.shape[1:])
+        else:
+            sample = jax.random.permutation(sampling_key, x)
+
         rng["jax_state"] = rng_key
         return (rng, sample)
 
@@ -321,7 +416,7 @@ def jax_sample_fn_permutation(op):
 
 
 @jax_sample_fn.register(ptr.BinomialRV)
-def jax_sample_fn_binomial(op):
+def jax_sample_fn_binomial(op, node):
     if not numpyro_available:
         raise NotImplementedError(
             f"No JAX implementation for the given distribution: {op.name}. "
@@ -344,7 +439,7 @@ def jax_sample_fn_binomial(op):
 
 
 @jax_sample_fn.register(ptr.MultinomialRV)
-def jax_sample_fn_multinomial(op):
+def jax_sample_fn_multinomial(op, node):
     if not numpyro_available:
         raise NotImplementedError(
             f"No JAX implementation for the given distribution: {op.name}. "
@@ -367,7 +462,7 @@ def jax_sample_fn_multinomial(op):
 
 
 @jax_sample_fn.register(ptr.VonMisesRV)
-def jax_sample_fn_vonmises(op):
+def jax_sample_fn_vonmises(op, node):
     if not numpyro_available:
         raise NotImplementedError(
             f"No JAX implementation for the given distribution: {op.name}. "

{pytensor-2.21.0 → pytensor-2.22.1}/pytensor/link/numba/dispatch/basic.py

@@ -604,36 +604,70 @@ def numba_funcify_IncSubtensor(op, node, **kwargs):
     return numba_njit(incsubtensor_fn, boundscheck=True)
 
 
-@numba_njit(boundscheck=True)
-def advancedincsubtensor1_inplace_set(x, vals, idxs):
-    for idx, val in zip(idxs, vals):
-        x[idx] = val
-    return x
-
-
-@numba_njit(boundscheck=True)
-def advancedincsubtensor1_inplace_inc(x, vals, idxs):
-    for idx, val in zip(idxs, vals):
-        x[idx] += val
-    return x
-
-
 @numba_funcify.register(AdvancedIncSubtensor1)
 def numba_funcify_AdvancedIncSubtensor1(op, node, **kwargs):
     inplace = op.inplace
     set_instead_of_inc = op.set_instead_of_inc
+    x, vals, idxs = node.inputs
+    # TODO: Add explicit expand_dims in make_node so we don't need to worry about this here
+    broadcast = vals.type.ndim < x.type.ndim or vals.type.broadcastable[0]
 
     if set_instead_of_inc:
-        advancedincsubtensor1_inplace = global_numba_func(
-            advancedincsubtensor1_inplace_set
-        )
+        if broadcast:
+
+            @numba_njit(boundscheck=True)
+            def advancedincsubtensor1_inplace(x, val, idxs):
+                if val.ndim == x.ndim:
+                    core_val = val[0]
+                elif val.ndim == 0:
+                    # Workaround for https://github.com/numba/numba/issues/9573
+                    core_val = val.item()
+                else:
+                    core_val = val
+
+                for idx in idxs:
+                    x[idx] = core_val
+                return x
+
+        else:
+
+            @numba_njit(boundscheck=True)
+            def advancedincsubtensor1_inplace(x, vals, idxs):
+                if not len(idxs) == len(vals):
+                    raise ValueError("The number of indices and values must match.")
+                for idx, val in zip(idxs, vals):
+                    x[idx] = val
+                return x
     else:
-        advancedincsubtensor1_inplace = global_numba_func(
-            advancedincsubtensor1_inplace_inc
-        )
+        if broadcast:
+
+            @numba_njit(boundscheck=True)
+            def advancedincsubtensor1_inplace(x, val, idxs):
+                if val.ndim == x.ndim:
+                    core_val = val[0]
+                elif val.ndim == 0:
+                    # Workaround for https://github.com/numba/numba/issues/9573
+                    core_val = val.item()
+                else:
+                    core_val = val
+
+                for idx in idxs:
+                    x[idx] += core_val
+                return x
+
+        else:
+
+            @numba_njit(boundscheck=True)
+            def advancedincsubtensor1_inplace(x, vals, idxs):
+                if not len(idxs) == len(vals):
+                    raise ValueError("The number of indices and values must match.")
+                for idx, val in zip(idxs, vals):
+                    x[idx] += val
+                return x
 
     if inplace:
-        return global_numba_func(advancedincsubtensor1_inplace)
+        return advancedincsubtensor1_inplace
+
     else:
 
         @numba_njit

{pytensor-2.21.0 → pytensor-2.22.1}/pytensor/link/numba/dispatch/random.py

@@ -210,7 +210,6 @@ def {sized_fn_name}({random_fn_input_names}):
 @numba_funcify.register(ptr.BinomialRV)
 @numba_funcify.register(ptr.MultinomialRV)
 @numba_funcify.register(ptr.RandIntRV)  # only the first two arguments are supported
-@numba_funcify.register(ptr.ChoiceRV)  # the `p` argument is not supported
 @numba_funcify.register(ptr.PermutationRV)
 def numba_funcify_RandomVariable(op, node, **kwargs):
     name = op.name
@@ -367,3 +366,63 @@ def numba_funcify_DirichletRV(op, node, **kwargs):
             return (rng, np.random.dirichlet(alphas, size))
 
     return dirichlet_rv
+
+
+@numba_funcify.register(ptr.ChoiceWithoutReplacement)
+def numba_funcify_choice_without_replacement(op, node, **kwargs):
+    batch_ndim = op.batch_ndim(node)
+    if batch_ndim:
+        # The code isn't too hard to write, but Numba doesn't support a with ndim > 1,
+        # and I don't want to change the batched tests for this
+        # We'll just raise an error for now
+        raise NotImplementedError(
+            "ChoiceWithoutReplacement with batch_ndim not supported in Numba backend"
+        )
+
+    [core_shape_len] = node.inputs[-1].type.shape
+
+    if op.has_p_param:
+
+        @numba_basic.numba_njit
+        def choice_without_replacement_rv(rng, size, dtype, a, p, core_shape):
+            core_shape = numba_ndarray.to_fixed_tuple(core_shape, core_shape_len)
+            samples = np.random.choice(a, size=core_shape, replace=False, p=p)
+            return (rng, samples)
+    else:
+
+        @numba_basic.numba_njit
+        def choice_without_replacement_rv(rng, size, dtype, a, core_shape):
+            core_shape = numba_ndarray.to_fixed_tuple(core_shape, core_shape_len)
+            samples = np.random.choice(a, size=core_shape, replace=False)
+            return (rng, samples)
+
+    return choice_without_replacement_rv
+
+
+@numba_funcify.register(ptr.PermutationRV)
+def numba_funcify_permutation(op, node, **kwargs):
+    # PyTensor uses size=() to represent size=None
+    size_is_none = node.inputs[1].type.shape == (0,)
+    batch_ndim = op.batch_ndim(node)
+    x_batch_ndim = node.inputs[-1].type.ndim - op.ndims_params[0]
+
+    @numba_basic.numba_njit
+    def permutation_rv(rng, size, dtype, x):
+        if batch_ndim:
+            x_core_shape = x.shape[x_batch_ndim:]
+            if size_is_none:
+                size = x.shape[:batch_ndim]
+            else:
+                size = numba_ndarray.to_fixed_tuple(size, batch_ndim)
+                x = np.broadcast_to(x, size + x_core_shape)
+
+            samples = np.empty(size + x_core_shape, dtype=x.dtype)
+            for index in np.ndindex(size):
+                samples[index] = np.random.permutation(x[index])
+
+        else:
+            samples = np.random.permutation(x)
+
+        return (rng, samples)
+
+    return permutation_rv

{pytensor-2.21.0 → pytensor-2.22.1}/pytensor/scalar/basic.py

@@ -24,7 +24,7 @@ import pytensor
 from pytensor import printing
 from pytensor.configdefaults import config
 from pytensor.gradient import DisconnectedType, grad_undefined
-from pytensor.graph.basic import Apply, Constant, Variable, clone, list_of_nodes
+from pytensor.graph.basic import Apply, Constant, Variable, applys_between, clone
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.op import HasInnerGraph
 from pytensor.graph.rewriting.basic import MergeOptimizer
@@ -4125,7 +4125,7 @@ class ScalarInnerGraphOp(ScalarOp, HasInnerGraph):
 
     def prepare_node(self, node, storage_map, compute_map, impl):
         if impl not in self.prepare_node_called:
-            for n in list_of_nodes(self.inputs, self.outputs):
+            for n in applys_between(self.inputs, self.outputs):
                 n.op.prepare_node(n, None, None, impl)
             self.prepare_node_called.add(impl)