pytensor 2.22.0__tar.gz → 2.22.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pytensor
-Version: 2.22.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.22.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.22.0 → pytensor-2.22.1}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-05-22T14:12:17+0200",
+ "date": "2024-05-28T12:10:34+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "bb028ae2330433755b9d4aa32ab6e8d0c9f662fc",
- "version": "2.22.0"
+ "full-revisionid": "fc21336a67334241d1c9a59383084ce453988d4e",
+ "version": "2.22.1"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.22.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]
 

{pytensor-2.22.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.22.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.22.0 → pytensor-2.22.1}/pytensor/tensor/rewriting/basic.py

@@ -351,10 +351,7 @@ def local_fill_sink(fgraph, node):
     # Check if we need to propagate the fill to the new outputs
     # It's enough to check the first output, as Elemwise outputs must all have the same shapes
     # Note: There are orderings that may require fewer fills.
-    old_bcast_pattern = node.outputs[0].type.broadcastable
-    models_iter = iter(models)
-    while old_bcast_pattern != outputs[0].type.broadcastable:
-        model = next(models_iter)
+    for model in models:
         # Only apply this model if it would actually do anything
         if broadcasted_by(outputs[0], model):
             outputs = [fill(model, output) for output in outputs]

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pytensor
-Version: 2.22.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: