CAReduce loop reordering C-impl

pytensor/tensor/math.py

 import builtins
 import warnings
 from collections.abc import Sequence
+from textwrap import dedent
 from typing import TYPE_CHECKING, Optional
 
 import numpy as np
@@ -361,12 +362,14 @@ class FixedOpCAReduce(CAReduce):
 
 
 class NonZeroDimsCAReduce(FixedOpCAReduce):
-    def _c_all(self, node, name, inames, onames, sub):
-        decl, checks, alloc, loop, end = super()._c_all(node, name, inames, onames, sub)
+    def _c_all(self, node, name, input_names, output_names, sub):
+        setup, alloc, loop, cast = super()._c_all(
+            node, name, input_names, output_names, sub
+        )
 
         # We add an additional check for zero-sized dimensions (This seems like
         # something that could enabled in `elemwise_cgen.make_checks`.)
-        iname = inames[0]
+        [iname] = input_names
 
         axis = self.axis
         if axis is None:
@@ -378,17 +381,19 @@ class NonZeroDimsCAReduce(FixedOpCAReduce):
 
         pattern_ = str(pattern)[1:-1]
 
-        decl += f"""int tosum[]={{{pattern_}}};"""
-        alloc += f"""
-                for(int i=0;i<PyArray_NDIM({iname});i++){{
-                    if(PyArray_DIMS({iname})[i]==0 && tosum[i]){{
-                        PyErr_Format(PyExc_ValueError,
-                            "Input of CAReduce{{{node.op.scalar_op}}} has zero-size on axis %%d",i);
-                        {sub["fail"]};
-                    }}
+        setup = f"int tosum[]={{{pattern_}}};" + setup
+        alloc += dedent(
+            f"""
+            for(int i=0;i<PyArray_NDIM({iname});i++){{
+                if(PyArray_DIMS({iname})[i]==0 && tosum[i]){{
+                    PyErr_Format(PyExc_ValueError,
+                        "Input of CAReduce{{{node.op.scalar_op}}} has zero-size on axis %%d",i);
+                    {sub["fail"]};
                 }}
-                """
-        return decl, checks, alloc, loop, end
+            }}
+            """
+        )
+        return setup, alloc, loop, cast
 
 
 class Max(NonZeroDimsCAReduce):