Reduce pytensor function call overhead

Author: ricardoV94

pytensor/compile/builders.py

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

pytensor/compile/function/types.py

@@ -87,10 +87,6 @@ def __set__(self, value: Any) -> None:
         if self.readonly:
             raise Exception(f"Cannot set readonly storage: {self.name}")
         try:
-            if value is None:
-                self.storage[0] = None
-                return
-
             kwargs = {}
             if self.strict:
                 kwargs["strict"] = True
@@ -539,14 +535,12 @@ def make_thunk(self, **kwargs):
 
         def f():
             for inputs in input_lists[1:]:
-                # strict=False because we are in a hot loop
-                for input1, input2 in zip(inputs0, inputs, strict=False):
+                for input1, input2 in zip(inputs0, inputs, strict=True):
                     input2.storage[0] = copy(input1.storage[0])
             for x in to_reset:
                 x[0] = None
             pre(self, [input.data for input in input_lists[0]], order, thunk_groups)
-            # strict=False because we are in a hot loop
-            for i, (thunks, node) in enumerate(zip(thunk_groups, order, strict=False)):
+            for i, (thunks, node) in enumerate(zip(thunk_groups, order, strict=True)):
                 try:
                     wrapper(self.fgraph, i, node, *thunks)
                 except Exception:
@@ -668,10 +662,12 @@ def thunk(
                 #  since the error may come from any of them?
                 raise_with_op(self.fgraph, output_nodes[0], thunk)
 
-            # strict=False because we are in a hot loop
-            for o_storage, o_val in zip(thunk_outputs, outputs, strict=False):
+            # strict=None because we are in a hot loop
+            for o_storage, o_val in zip(thunk_outputs, outputs):  # noqa: B905
                 o_storage[0] = o_val
 
+            return outputs
+
         thunk.inputs = thunk_inputs
         thunk.outputs = thunk_outputs
         thunk.lazy = False

pytensor/link/basic.py

@@ -1988,26 +1988,25 @@ def make_thunk(self, **kwargs):
         )
 
         def f():
-            # strict=False because we are in a hot loop
-            for input1, input2 in zip(i1, i2, strict=False):
+            for input1, input2 in zip(i1, i2, strict=True):
                 # 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=False
+                thunks1, thunks2, order1, order2, strict=True
             ):
-                for output, storage in zip(node1.outputs, thunk1.outputs, strict=False):
+                for output, storage in zip(node1.outputs, thunk1.outputs, strict=True):
                     if output in no_recycling:
                         storage[0] = None
-                for output, storage in zip(node2.outputs, thunk2.outputs, strict=False):
+                for output, storage in zip(node2.outputs, thunk2.outputs, strict=True):
                     if output in no_recycling:
                         storage[0] = None
                 try:
                     thunk1()
                     thunk2()
                     for output1, output2 in zip(
-                        thunk1.outputs, thunk2.outputs, strict=False
+                        thunk1.outputs, thunk2.outputs, strict=True
                     ):
                         self.checker(output1, output2)
                 except Exception:

pytensor/link/c/basic.py

@@ -312,10 +312,10 @@ def py_perform_return(inputs):
     else:
 
         def py_perform_return(inputs):
-            # strict=False because we are in a hot loop
+            # strict=None because we are in a hot loop
             return tuple(
                 out_type.filter(out[0])
-                for out_type, out in zip(output_types, py_perform(inputs), strict=False)
+                for out_type, out in zip(output_types, py_perform(inputs))  # noqa: B905
             )
 
     @numba_njit

pytensor/link/numba/dispatch/basic.py

@@ -207,8 +207,8 @@ def streamline_nice_errors_f():
             for x in no_recycling:
                 x[0] = None
             try:
-                # strict=False because we are in a hot loop
-                for thunk, node in zip(thunks, order, strict=False):
+                # strict=None because we are in a hot loop
+                for thunk, node in zip(thunks, order):  # noqa: B905
                     thunk()
             except Exception:
                 raise_with_op(fgraph, node, thunk)

pytensor/link/utils.py

@@ -4427,8 +4427,8 @@ def make_node(self, *inputs):
 
     def perform(self, node, inputs, output_storage):
         outputs = self.py_perform_fn(*inputs)
-        # strict=False because we are in a hot loop
-        for storage, out_val in zip(output_storage, outputs, strict=False):
+        # strict=None because we are in a hot loop
+        for storage, out_val in zip(output_storage, outputs):  # noqa: B905
             storage[0] = out_val
 
     def grad(self, inputs, output_grads):

pytensor/scalar/basic.py

@@ -207,8 +207,8 @@ def perform(self, node, inputs, output_storage):
             for i in range(n_steps):
                 carry = inner_fn(*carry, *constant)
 
-        # strict=False because we are in a hot loop
-        for storage, out_val in zip(output_storage, carry, strict=False):
+        # strict=None because we are in a hot loop
+        for storage, out_val in zip(output_storage, carry):  # noqa: B905
             storage[0] = out_val
 
     @property

pytensor/scalar/loop.py

@@ -1865,8 +1865,8 @@ def rng_fn(cls, rng, p, size):
             # to `p.shape[:-1]` in the call to `vsearchsorted` below.
             if len(size) < (p.ndim - 1):
                 raise ValueError("`size` is incompatible with the shape of `p`")
-            # strict=False because we are in a hot loop
-            for s, ps in zip(reversed(size), reversed(p.shape[:-1]), strict=False):
+            # strict=None because we are in a hot loop
+            for s, ps in zip(reversed(size), reversed(p.shape[:-1])):  # noqa: B905
                 if s == 1 and ps != 1:
                     raise ValueError("`size` is incompatible with the shape of `p`")
 

pytensor/tensor/random/basic.py

@@ -44,8 +44,8 @@ def params_broadcast_shapes(
     max_fn = maximum if use_pytensor else max
 
     rev_extra_dims: list[int] = []
-    # strict=False because we are in a hot loop
-    for ndim_param, param_shape in zip(ndims_params, param_shapes, strict=False):
+    # strict=None because we are in a hot loop
+    for ndim_param, param_shape in zip(ndims_params, param_shapes):  # noqa: B905
         # We need this in order to use `len`
         param_shape = tuple(param_shape)
         extras = tuple(param_shape[: (len(param_shape) - ndim_param)])
@@ -69,7 +69,7 @@ def max_bcast(x, y):
         (extra_dims + tuple(param_shape)[-ndim_param:])
         if ndim_param > 0
         else extra_dims
-        for ndim_param, param_shape in zip(ndims_params, param_shapes, strict=False)
+        for ndim_param, param_shape in zip(ndims_params, param_shapes)  # noqa: B905
     ]
 
     return bcast_shapes
@@ -127,10 +127,10 @@ def broadcast_params(
     )
     broadcast_to_fn = broadcast_to if use_pytensor else np.broadcast_to
 
-    # strict=False because we are in a hot loop
+    # strict=None because we are in a hot loop
     bcast_params = [
         broadcast_to_fn(param, shape)
-        for shape, param in zip(shapes, params, strict=False)
+        for shape, param in zip(shapes, params)  # noqa: B905
     ]
 
     return bcast_params