Add exceptions for hot loops

Author: Armavica

pytensor/compile/builders.py

@@ -863,5 +863,6 @@ def clone(self):
     def perform(self, node, inputs, outputs):
         variables = self.fn(*inputs)
         assert len(variables) == len(outputs)
-        for output, variable in zip(outputs, variables, strict=True):
+        # strict=False because asserted above
+        for output, variable in zip(outputs, variables, strict=False):
             output[0] = variable

pytensor/compile/function/types.py

@@ -1002,8 +1002,9 @@ def __call__(self, *args, **kwargs):
         # if we are allowing garbage collection, remove the
         # output reference from the internal storage cells
         if getattr(self.vm, "allow_gc", False):
+            # strict=False because we are in a hot loop
             for o_container, o_variable in zip(
-                self.output_storage, self.maker.fgraph.outputs, strict=True
+                self.output_storage, self.maker.fgraph.outputs, strict=False
             ):
                 if o_variable.owner is not None:
                     # this node is the variable of computation

pytensor/ifelse.py

@@ -305,7 +305,8 @@ def thunk():
                     if len(ls) > 0:
                         return ls
                     else:
-                        for out, t in zip(outputs, input_true_branch, strict=True):
+                        # strict=False because we are in a hot loop
+                        for out, t in zip(outputs, input_true_branch, strict=False):
                             compute_map[out][0] = 1
                             val = storage_map[t][0]
                             if self.as_view:
@@ -325,7 +326,8 @@ def thunk():
                     if len(ls) > 0:
                         return ls
                     else:
-                        for out, f in zip(outputs, inputs_false_branch, strict=True):
+                        # strict=False because we are in a hot loop
+                        for out, f in zip(outputs, inputs_false_branch, strict=False):
                             compute_map[out][0] = 1
                             # can't view both outputs unless destroyhandler
                             # improves

pytensor/link/basic.py

@@ -544,7 +544,8 @@ def f():
             for x in to_reset:
                 x[0] = None
             pre(self, [input.data for input in input_lists[0]], order, thunk_groups)
-            for i, (thunks, node) in enumerate(zip(thunk_groups, order, strict=True)):
+            # strict=False because we are in a hot loop
+            for i, (thunks, node) in enumerate(zip(thunk_groups, order, strict=False)):
                 try:
                     wrapper(self.fgraph, i, node, *thunks)
                 except Exception:
@@ -666,8 +667,9 @@ def thunk(
         ):
             outputs = fgraph_jit(*[self.input_filter(x[0]) for x in thunk_inputs])
 
+            # strict=False because we are in a hot loop
             for o_var, o_storage, o_val in zip(
-                fgraph.outputs, thunk_outputs, outputs, strict=True
+                fgraph.outputs, thunk_outputs, outputs, strict=False
             ):
                 compute_map[o_var][0] = True
                 o_storage[0] = self.output_filter(o_var, o_val)

pytensor/link/c/basic.py

@@ -1993,25 +1993,26 @@ def make_thunk(self, **kwargs):
         )
 
         def f():
-            for input1, input2 in zip(i1, i2, strict=True):
+            # strict=False because we are in a hot loop
+            for input1, input2 in zip(i1, i2, strict=False):
                 # Set the inputs to be the same in both branches.
                 # The copy is necessary in order for inplace ops not to
                 # interfere.
                 input2.storage[0] = copy(input1.storage[0])
             for thunk1, thunk2, node1, node2 in zip(
-                thunks1, thunks2, order1, order2, strict=True
+                thunks1, thunks2, order1, order2, strict=False
             ):
-                for output, storage in zip(node1.outputs, thunk1.outputs, strict=True):
+                for output, storage in zip(node1.outputs, thunk1.outputs, strict=False):
                     if output in no_recycling:
                         storage[0] = None
-                for output, storage in zip(node2.outputs, thunk2.outputs, strict=True):
+                for output, storage in zip(node2.outputs, thunk2.outputs, strict=False):
                     if output in no_recycling:
                         storage[0] = None
                 try:
                     thunk1()
                     thunk2()
                     for output1, output2 in zip(
-                        thunk1.outputs, thunk2.outputs, strict=True
+                        thunk1.outputs, thunk2.outputs, strict=False
                     ):
                         self.checker(output1, output2)
                 except Exception:

pytensor/link/numba/dispatch/subtensor.py

@@ -158,7 +158,9 @@ def advancedincsubtensor1_inplace(x, val, idxs):
             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, strict=True):
+                # strict=False to not add overhead in the jitted code
+                # TODO: check
+                for idx, val in zip(idxs, vals, strict=False):
                     x[idx] = val
                 return x
     else:

pytensor/link/pytorch/dispatch/shape.py

@@ -34,7 +34,8 @@ def shape_i(x):
 def pytorch_funcify_SpecifyShape(op, node, **kwargs):
     def specifyshape(x, *shape):
         assert x.ndim == len(shape)
-        for actual, expected in zip(x.shape, shape, strict=True):
+        # strict=False because asserted above
+        for actual, expected in zip(x.shape, shape, strict=False):
             if expected is None:
                 continue
             if actual != expected:

pytensor/scalar/basic.py

@@ -1150,8 +1150,9 @@ def perform(self, node, inputs, output_storage):
         else:
             variables = from_return_values(self.impl(*inputs))
             assert len(variables) == len(output_storage)
+            # strict=False because we are in a hot loop
             for out, storage, variable in zip(
-                node.outputs, output_storage, variables, strict=True
+                node.outputs, output_storage, variables, strict=False
             ):
                 dtype = out.dtype
                 storage[0] = self._cast_scalar(variable, dtype)
@@ -4328,7 +4329,8 @@ def make_node(self, *inputs):
 
     def perform(self, node, inputs, output_storage):
         outputs = self.py_perform_fn(*inputs)
-        for storage, out_val in zip(output_storage, outputs, strict=True):
+        # strict=False because we are in a hot loop
+        for storage, out_val in zip(output_storage, outputs, strict=False):
             storage[0] = out_val
 
     def grad(self, inputs, output_grads):

pytensor/scalar/loop.py

@@ -93,7 +93,7 @@ def _validate_updates(
                 )
         else:
             update = outputs
-        for i, u in zip(init[: len(update)], update, strict=True):
+        for i, u in zip(init, update, strict=False):
             if i.type != u.type:
                 raise TypeError(
                     "Init and update types must be the same: "
@@ -207,7 +207,8 @@ def perform(self, node, inputs, output_storage):
             for i in range(n_steps):
                 carry = inner_fn(*carry, *constant)
 
-        for storage, out_val in zip(output_storage, carry, strict=True):
+        # strict=False because we are in a hot loop
+        for storage, out_val in zip(output_storage, carry, strict=False):
             storage[0] = out_val
 
     @property

pytensor/tensor/rewriting/subtensor.py

@@ -683,7 +683,7 @@ def local_subtensor_of_alloc(fgraph, node):
     # Slices to take from val
     val_slices = []
 
-    for i, (sl, dim) in enumerate(zip(slices, dims[: len(slices)], strict=True)):
+    for i, (sl, dim) in enumerate(zip(slices, dims, strict=False)):
         # If val was not copied over that dim,
         # we need to take the appropriate subtensor on it.
         if i >= n_added_dims: