Specialize Zero Alloc

Author: ricardoV94

pytensor/tensor/basic.py

@@ -1634,6 +1634,14 @@ def _check_runtime_broadcast(node, value, shape):
             if v_static_dim is None and value_dim == 1 and out_dim != 1:
                 raise ValueError(Alloc._runtime_broadcast_error_msg)
 
+    @staticmethod
+    def value_is_scalar_zero(x: TensorVariable) -> bool:
+        return (
+            all(x.type.broadcastable)
+            and isinstance(x, Constant)
+            and (x.unique_value == 0)
+        )
+
     def perform(self, node, inputs, out_):
         (out,) = out_
         v = inputs[0]
@@ -1659,6 +1667,7 @@ def c_code(self, node, name, inp, out, sub):
         o_static_shape = node.outputs[0].type.shape
         v_ndim = len(v_static_shape)
         o_ndim = len(o_static_shape)
+        is_zero = self.value_is_scalar_zero(node.inputs[0])
         assert o_ndim == len(inp[1:])
 
         # Declare variables
@@ -1699,16 +1708,18 @@ def c_code(self, node, name, inp, out, sub):
                     {fail}
                 }}
             }}
-
+            if ({int(is_zero)} && (PyArray_IS_C_CONTIGUOUS({zz}) || PyArray_IS_F_CONTIGUOUS({zz}))){{
+                PyArray_FILLWBYTE({zz}, 0);
+            }}
             // This function takes care of broadcasting
-            if (PyArray_CopyInto({zz}, {vv}) == -1)
+            else if (PyArray_CopyInto({zz}, {vv}) == -1)
               {fail}
             """
 
         return code
 
     def c_code_cache_version(self):
-        return (4,)
+        return (5,)
 
     def infer_shape(self, fgraph, node, input_shapes):
         return [node.inputs[1:]]