Format kron atom and mirror left_matmul structure

Apply clang-format and tidy the native kron atom: factor parameter
refresh into a helper, null freed pointers, and trim doc duplicated by
the kron_expr definition in subexpr.h.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: Transurgeon

include/atoms/affine.h

@@ -85,7 +85,7 @@ expr *new_convolve(expr *param_node, expr *child);
    variables and is passed as child. const_is_left selects which operand is the
    parameter: 1 -> A=param_node, B=child; 0 -> A=child, B=param_node. (p, q) are
    A's dims and (r, s) are B's dims. */
-expr *new_kron(expr *param_node, expr *child, int const_is_left, int p, int q,
-               int r, int s);
+expr *new_kron(expr *param_node, expr *child, int const_is_left, int p, int q, int r,
+               int s);
 
 #endif /* AFFINE_H */

src/atoms/affine/kron.c

@@ -24,8 +24,8 @@
 #include <stdlib.h>
 #include <string.h>
 
-/* Kronecker product Z = kron(A, B), where exactly one operand is variable-free
- * (param_source) and the other (child = node->left) carries the variables.
+/* Kronecker product Z = kron(A, B). See the kron_expr definition in subexpr.h
+ * for the operand layout; the index math behind the scaled gather is below.
  *
  * With column-major (Fortran) flattening, an output index OUT = I + J*(p*r)
  * decomposes as I = i*r + k and J = j*s + l (i in [0,p), k in [0,r), j in [0,q),
@@ -40,18 +40,24 @@
  * forward, Jacobian and (affine) Hessian are all scaled gathers -- no
  * size_out x size_child coefficient matrix and no sparse matmul. */
 
+/* Pull current parameter values through any broadcast/promote wrappers. */
+static void refresh_param_values(kron_expr *knode)
+{
+    if (!knode->base.needs_parameter_refresh)
+    {
+        return;
+    }
+
+    knode->param_source->forward(knode->param_source, NULL);
+    knode->base.needs_parameter_refresh = false;
+}
+
 static void forward(expr *node, const double *u)
 {
     expr *child = node->left;
     kron_expr *knode = (kron_expr *) node;
 
-    /* Pull current parameter values through any broadcast/promote wrappers. */
-    if (knode->base.needs_parameter_refresh)
-    {
-        knode->param_source->forward(knode->param_source, NULL);
-        knode->base.needs_parameter_refresh = false;
-    }
-
+    refresh_param_values(knode);
     child->forward(child, u);
 
     const double *a = knode->param_source->value;
@@ -163,10 +169,14 @@ static void free_type_data(expr *node)
     sp_free(knode->child_row);
     sp_free(knode->coeff_idx);
     free_expr(knode->param_source);
+
+    knode->child_row = NULL;
+    knode->coeff_idx = NULL;
+    knode->param_source = NULL;
 }
 
-expr *new_kron(expr *param_node, expr *child, int const_is_left, int p, int q,
-               int r, int s)
+expr *new_kron(expr *param_node, expr *child, int const_is_left, int p, int q, int r,
+               int s)
 {
     int d1 = p * r;
     int d2 = q * s;

tests/all_tests.c

@@ -8,9 +8,9 @@
 #include "forward_pass/affine/test_add.h"
 #include "forward_pass/affine/test_broadcast.h"
 #include "forward_pass/affine/test_convolve.h"
-#include "forward_pass/affine/test_kron.h"
 #include "forward_pass/affine/test_diag_mat.h"
 #include "forward_pass/affine/test_hstack.h"
+#include "forward_pass/affine/test_kron.h"
 #include "forward_pass/affine/test_left_matmul_dense.h"
 #include "forward_pass/affine/test_linear_op.h"
 #include "forward_pass/affine/test_neg.h"
@@ -28,10 +28,10 @@
 #include "forward_pass/other/test_prod_axis_zero.h"
 #include "jacobian_tests/affine/test_broadcast.h"
 #include "jacobian_tests/affine/test_convolve.h"
-#include "jacobian_tests/affine/test_kron.h"
 #include "jacobian_tests/affine/test_diag_mat.h"
 #include "jacobian_tests/affine/test_hstack.h"
 #include "jacobian_tests/affine/test_index.h"
+#include "jacobian_tests/affine/test_kron.h"
 #include "jacobian_tests/affine/test_left_matmul.h"
 #include "jacobian_tests/affine/test_neg.h"
 #include "jacobian_tests/affine/test_promote.h"
@@ -75,10 +75,10 @@
 #include "utils/test_stacked_pd.h"
 #include "wsum_hess/affine/test_broadcast.h"
 #include "wsum_hess/affine/test_convolve.h"
-#include "wsum_hess/affine/test_kron.h"
 #include "wsum_hess/affine/test_diag_mat.h"
 #include "wsum_hess/affine/test_hstack.h"
 #include "wsum_hess/affine/test_index.h"
+#include "wsum_hess/affine/test_kron.h"
 #include "wsum_hess/affine/test_left_matmul.h"
 #include "wsum_hess/affine/test_right_matmul.h"
 #include "wsum_hess/affine/test_scalar_mult.h"

tests/forward_pass/affine/test_kron.h

@@ -42,8 +42,8 @@ const char *test_kron_forward_const_right(void)
     double u[4] = {5.0, 7.0, 6.0, 8.0}; /* col-major [[5,6],[7,8]] */
     Z->forward(Z, u);
 
-    double expected[16] = {5,  15, 7,  21, 10, 20, 14, 28,
-                           6,  18, 8,  24, 12, 24, 16, 32};
+    double expected[16] = {5, 15, 7, 21, 10, 20, 14, 28,
+                           6, 18, 8, 24, 12, 24, 16, 32};
 
     mu_assert("kron const-right d1=4", Z->d1 == 4);
     mu_assert("kron const-right d2=4", Z->d2 == 4);
@@ -70,8 +70,7 @@ const char *test_kron_forward_scalar(void)
 
     mu_assert("kron scalar d1=2", Z->d1 == 2);
     mu_assert("kron scalar d2=2", Z->d2 == 2);
-    mu_assert("kron scalar forward failed",
-              cmp_double_array(Z->value, expected, 4));
+    mu_assert("kron scalar forward failed", cmp_double_array(Z->value, expected, 4));
 
     free_expr(Z);
     return 0;

tests/jacobian_tests/affine/test_kron.h

@@ -27,15 +27,13 @@ const char *test_jacobian_kron_const_left(void)
     mu_assert("kron J cols", Z->jacobian->n == 4);
     mu_assert("kron J nnz", Z->jacobian->nnz == 16);
 
-    int expected_p[17] = {0, 1, 2, 3, 4, 5, 6, 7, 8,
-                          9, 10, 11, 12, 13, 14, 15, 16};
+    int expected_p[17] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16};
     int expected_i[16] = {0, 1, 0, 1, 2, 3, 2, 3, 0, 1, 0, 1, 2, 3, 2, 3};
     double expected_x[16] = {1, 1, 3, 3, 1, 1, 3, 3, 2, 2, 4, 4, 2, 2, 4, 4};
 
     mu_assert("kron const-left J sparsity",
               cmp_sparsity(Z->jacobian, expected_p, expected_i, 16, 16));
-    mu_assert("kron const-left J values",
-              cmp_values(Z->jacobian, expected_x, 16));
+    mu_assert("kron const-left J values", cmp_values(Z->jacobian, expected_x, 16));
 
     free_expr(Z);
     return 0;
@@ -57,15 +55,13 @@ const char *test_jacobian_kron_const_right(void)
     mu_assert("kron J rows", Z->jacobian->m == 16);
     mu_assert("kron J nnz", Z->jacobian->nnz == 16);
 
-    int expected_p[17] = {0, 1, 2, 3, 4, 5, 6, 7, 8,
-                          9, 10, 11, 12, 13, 14, 15, 16};
+    int expected_p[17] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16};
     int expected_i[16] = {0, 0, 1, 1, 0, 0, 1, 1, 2, 2, 3, 3, 2, 2, 3, 3};
     double expected_x[16] = {1, 3, 1, 3, 2, 4, 2, 4, 1, 3, 1, 3, 2, 4, 2, 4};
 
     mu_assert("kron const-right J sparsity",
               cmp_sparsity(Z->jacobian, expected_p, expected_i, 16, 16));
-    mu_assert("kron const-right J values",
-              cmp_values(Z->jacobian, expected_x, 16));
+    mu_assert("kron const-right J values", cmp_values(Z->jacobian, expected_x, 16));
 
     free_expr(Z);
     return 0;