pytensor 3.2.0__tar.gz → 3.2.1__tar.gz

{pytensor-3.2.0/pytensor.egg-info → pytensor-3.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 3.2.0
+Version: 3.2.1
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License-Expression: BSD-3-Clause

{pytensor-3.2.0 → pytensor-3.2.1}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-07-16T16:05:00+0200",
+ "date": "2026-07-17T16:59:23+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "5480dc534c09f11c8cc5b0cf66374eabb44f0c6e",
- "version": "3.2.0"
+ "full-revisionid": "491faffa05d1e92419dba8e02b245eb2b6227a60",
+ "version": "3.2.1"
 }
 '''  # END VERSION_JSON
 

{pytensor-3.2.0 → pytensor-3.2.1}/pytensor/link/jax/dispatch/tensor_basic.py

@@ -71,7 +71,10 @@ def jax_funcify_ARange(op, node, **kwargs):
             case (Subtensor(), shape_var, *_) if isinstance(shape_var.owner_op, Shape):
                 constant_args.append(None)
             case _ if isinstance(arg, Constant):
-                constant_args.append(arg.value)
+                # Cast to the Op's dtype: PyTensor types integer literals (e.g. 0/1 arange
+                # start/step) as int8, and jnp.arange bounds-checks stop against the argument dtype,
+                # overflowing for stop > 127.
+                constant_args.append(np.asarray(arg.value, op.dtype))
             case _:
                 # TODO: This might be failing without need (e.g., if arg = shape(x)[-1] + 1)!
                 raise NotImplementedError(ARANGE_CONCRETE_VALUE_ERROR)

{pytensor-3.2.0 → pytensor-3.2.1}/pytensor/tensor/rewriting/subtensor.py

@@ -76,7 +76,6 @@ from pytensor.tensor.subtensor import (
     _is_provably_non_negative,
     _is_provably_positive,
     _non_consecutive_adv_indexing,
-    advanced_inc_subtensor,
     as_index_literal,
     flatten_index_variables,
     get_canonical_form_slice,
@@ -1104,11 +1103,13 @@ def local_useless_inc_subtensor(fgraph, node):
 @register_specialize
 @node_rewriter([AdvancedIncSubtensor])
 def local_set_to_inc_subtensor(fgraph, node):
-    r"""
-    AdvancedIncSubtensor(x, x[ilist]+other, ilist, set_instead_of_inc=True) ->
-    AdvancedIncSubtensor(x, other, ilist, set_instead_of_inc=False)
+    r"""Set of a read at the same indices: ``x[idx].set(x[idx] + other)``.
 
-    Only valid when ``ilist`` is duplicate-free: a dense set is last-wins while an
+    .. code::
+
+        x[idx].set(x[idx] + other) -> x[idx].inc(other)
+
+    Only valid when ``idx`` is duplicate-free: a dense set is last-wins while an
     inc accumulates, so duplicate indices would over-count.
 
     TODO FIXME: Why doesn't this apply to all `*IncSubtensor*` `Op`\s?  If it
@@ -1136,14 +1137,30 @@ def local_set_to_inc_subtensor(fgraph, node):
         other = addn.inputs[0]
     else:
         return
-    if subn.inputs[1] != node.inputs[2] or subn.inputs[0] != node.inputs[0]:
+    # The read has to gather exactly the positions the write stores to: same base,
+    # same axes, same indices. `idx_list` must be compared too -- it is what pins
+    # the indices to their axes, so without it a read of ``x[:, i]`` would match a
+    # write to ``x[i]``.
+    if (
+        subn.inputs[0] != node.inputs[0]
+        or subn.op.idx_list != node.op.idx_list
+        or subn.inputs[1:] != node.inputs[2:]
+    ):
         return
-    # set->inc is only valid when ilist is duplicate-free: ``set(x[ilist] + other)``
-    # is last-wins at repeated positions, while ``inc(other)`` would accumulate the
-    # contributions of every occurrence and over-count them.
-    if not _index_provably_unique(node.inputs[2], fgraph):
+    # set->inc is only valid when the written positions are duplicate-free:
+    # ``set(x[idx] + other)`` is last-wins at repeated positions, while
+    # ``inc(other)`` would accumulate the contributions of every occurrence and
+    # over-count them.
+    indices = indices_from_subtensor(node.inputs[2:], node.op.idx_list)
+    if not _advanced_indices_jointly_unique(indices, fgraph):
         return
-    ret = advanced_inc_subtensor(node.inputs[0], other, node.inputs[2])
+    new_op = type(node.op)(
+        idx_list=node.op.idx_list,
+        inplace=node.op.inplace,
+        set_instead_of_inc=False,
+        ignore_duplicates=node.op.ignore_duplicates,
+    )
+    ret = new_op(node.inputs[0], other, *node.inputs[2:])
 
     copy_stack_trace(node.outputs, ret)
 
@@ -1341,52 +1358,6 @@ def local_convert_negative_indices(fgraph, node):
     return [new_subtensor]
 
 
-@register_canonicalize
-@register_specialize
-@node_rewriter([AdvancedSubtensor])
-def local_useless_AdvancedSubtensor1(fgraph, node):
-    """Remove a leading-axis vector take that selects the full input.
-
-    The full input is taken when the index is equivalent to
-    ``arange(0, input.shape[0], 1)`` using either an explicit list/vector or
-    the `ARange` `Op`.
-
-    """
-    if node.op.idx_list != (0,):
-        # Only the leading integer-vector take (x[ilist])
-        return
-
-    # This optimization needs ShapeOpt and fgraph.shape_feature
-    if not hasattr(fgraph, "shape_feature"):
-        return
-
-    shape_feature = fgraph.shape_feature
-
-    # get length of the indexed tensor along the first axis
-    try:
-        length = get_scalar_constant_value(
-            shape_feature.get_shape(node.inputs[0], 0), only_process_constants=True
-        )
-    except NotScalarConstantError:
-        return False
-
-    # get index (which must be a vector by definition)
-    idx = node.inputs[1]
-
-    # `idx` must be equivalent to [0,1,...,shape[0] - 1] to qualify for
-    # this optimization
-    if isinstance(idx, Constant):
-        idx = idx.value
-        if len(idx) != length:
-            return False
-        if np.any(idx != np.arange(length)):
-            return False
-    else:
-        return False
-
-    return [node.inputs[0]]
-
-
 def _arange_index_to_slice(idx):
     """Return the Python ``slice`` equivalent to ``idx`` if it is a constant or
     symbolic arange with non-negative offset, else ``None``.
@@ -2343,8 +2314,7 @@ def local_write_of_write_same_indices(fgraph, node):
     base, a, *inner_idx_vars = inner_x.owner.inputs
 
     # Same indices: idx_list (slice specs) must match and all index
-    # variables must be identical.  AdvancedIncSubtensor1 has a fixed
-    # class-level idx_list = (0,) so the comparison is trivially true.
+    # variables must be identical.
     if node.op.idx_list != inner_x.owner.op.idx_list:
         return
     if not all(o is i for o, i in zip(outer_idx_vars, inner_idx_vars, strict=True)):

{pytensor-3.2.0 → pytensor-3.2.1}/pytensor/tensor/rewriting/subtensor_lift.py

@@ -1062,6 +1062,11 @@ def local_subtensor_make_vector(fgraph, node):
             except NotScalarConstantError:
                 pass
         elif idx.ndim == 1 and isinstance(idx, Constant):
+            if idx.type.dtype == "bool":
+                # A boolean mask selects the entries it is True at, which is not
+                # what mapping it to positions below would read it as. Leave it
+                # for `bool_idx_to_nonzero` to turn into integer positions.
+                return False
             values = list(map(int, list(idx.value)))
             if len(values) == 1:
                 ret = expand_dims(x.owner.inputs[values[0]], axis=0)

{pytensor-3.2.0 → pytensor-3.2.1/pytensor.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 3.2.0
+Version: 3.2.1
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License-Expression: BSD-3-Clause