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,8 +381,9 @@ class NonZeroDimsCAReduce(FixedOpCAReduce):
 
         pattern_ = str(pattern)[1:-1]
 
-        decl += f"""int tosum[]={{{pattern_}}};"""
-        alloc += f"""
+        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,
@@ -388,7 +392,8 @@ class NonZeroDimsCAReduce(FixedOpCAReduce):
                 }}
             }}
             """
-        return decl, checks, alloc, loop, end
+        )
+        return setup, alloc, loop, cast
 
 
 class Max(NonZeroDimsCAReduce):