clean up

src/atoms/bivariate_full_dom/matmul.c

@@ -214,8 +214,7 @@ static void eval_jacobian_no_chain_rule(expr *node)
 //  Jacobian initialization and evaluation for the case where chain rule is needed,
 //  ie., when at least one child is composite, or same variable. The jacobian of h(u)
 //  = f(u) @ g(u) where f is m x k and g is k x n, is given by
-//      Jh = (g^T kron I_m) Jf + (I_n kron f) Jg,
-//  the two kron dispatchers with g^T and f held as full permuted_dense operands.
+//      Jh = (g^T kron I_m) Jf + (I_n kron f) Jg.
 // ------------------------------------------------------------------------------------
 static void jacobian_init_chain_rule(expr *node)
 {
@@ -254,17 +253,13 @@ static void eval_jacobian_chain_rule(expr *node)
     eval_jacobian(f);
     eval_jacobian(g);
 
-    /* g->value is Y (k x n) column-major, which is Y^T row-major; f->value is X
-       (m x k) column-major, transposed into X's row-major block. */
+    /* refresh the values of Y^T and X */
     memcpy(mnode->YT->x, g->value, (size_t) k * n * sizeof(double));
     matrix_values_changed(mnode->YT);
     A_transpose(mnode->X->x, f->value, k, m);
     matrix_values_changed(mnode->X);
 
-    /* evaluate term1, term2, and their sum. term1 is stacked_pd and term2 is
-       stacked_pd for n > 1 (permuted_dense for n == 1); the sum reads them
-       through their CSR views, one O(nnz) refresh per stacked_pd term per
-       eval; a bound-map sum in the library would remove that. */
+    /* Jh = term1 + term2 */
     BA_dense_kron_I_matrices_fill_values(mnode->YT, m, f->jacobian, mnode->term1);
     BA_kron_matrices_fill_values(mnode->X, n, g->jacobian, mnode->term2);
     sum_matrices_fill_values(mnode->term1, mnode->term2, node->jacobian);

src/atoms/bivariate_full_dom/multiply.c

@@ -29,8 +29,6 @@
 
 // ------------------------------------------------------------------------------
 // Implementation of elementwise multiplication when both arguments are vectors.
-// If one argument is a scalar variable, the broadcasting should be represented
-// as a linear operator child node? How to treat if one variable is a constant?
 // ------------------------------------------------------------------------------
 static void forward(expr *node, const double *u)
 {
@@ -162,8 +160,7 @@ static void wsum_hess_init_impl(expr *node)
         mul_node->C = C;
         mul_node->CT = CT;
 
-        /* sparsity pattern of H = C + C^T + term2 + term3; the sum binds the
-           per-operand index maps it accumulates through at eval */
+        /* sparsity pattern of H = C + C^T + term2 + term3 */
         node->wsum_hess = sum_4_matrices_alloc(C, CT, x->wsum_hess, y->wsum_hess);
     }
 }

src/atoms/bivariate_restricted_dom/rel_entr.c

@@ -206,6 +206,14 @@ expr *new_rel_entr_vector_args(expr *left, expr *right)
         return NULL;
     }
 
+    /* the derivative code assumes two distinct variables */
+    if (left->var_id == right->var_id)
+    {
+        fprintf(stderr, "Error: Arguments of relative entropy must be distinct "
+                        "variables.\n");
+        return NULL;
+    }
+
     expr *node = (expr *) sp_calloc(1, sizeof(expr));
     init_expr(node, left->d1, left->d2, left->n_vars, forward_vector_args,
               jacobian_init_vectors_args, eval_jacobian_vector_args, is_affine,

src/atoms/elementwise_full_dom/common.c

@@ -67,9 +67,8 @@ void eval_jacobian_elementwise(expr *node)
         /* jacobian of h(x) = f(g(x)) is Jf @ Jg, and here Jf is diagonal */
         eval_jacobian(child);
         node->local_jacobian(node, node->work->local_jac_diag);
-        memcpy(node->work->dwork, node->work->local_jac_diag,
-               node->size * sizeof(double));
-        DA_matrices_fill_values(node->work->dwork, child->jacobian, node->jacobian);
+        DA_matrices_fill_values(node->work->local_jac_diag, child->jacobian,
+                                node->jacobian);
     }
 }
 

src/atoms/elementwise_full_dom/logistic.c

@@ -67,7 +67,8 @@ static void local_wsum_hess(expr *node, double *out, const double *w)
     }
     else
     {
-        sigmas = node->work->dwork;
+        /* local_jac_diag holds sigma(g(x)) from eval_jacobian */
+        sigmas = node->work->local_jac_diag;
     }
 
     for (int j = 0; j < node->size; j++)

tests/forward_pass/elementwise_restricted_dom/test_log.h

@@ -33,6 +33,8 @@ const char *test_restricted_rejects_non_variable_child(void)
     mu_assert("log of affine child should be NULL", new_log(Ax) == NULL);
     mu_assert("rel_entr of affine child should be NULL",
               new_rel_entr_vector_args(Ax, x) == NULL);
+    mu_assert("rel_entr of same variable should be NULL",
+              new_rel_entr_vector_args(x, x) == NULL);
     mu_assert("quad_over_lin with non-scalar affine denominator should be NULL",
               new_quad_over_lin(x, Ax) == NULL);