add initial attempt for fixing parameters and broadcasting

Author: Transurgeon

src/atoms/affine/left_matmul.c

@@ -63,6 +63,14 @@ static void refresh_param_values(left_matmul_expr *lnode)
 static void forward(expr *node, const double *u)
 {
     left_matmul_expr *lnode = (left_matmul_expr *) node;
+
+    /* Refresh param_source expression tree if parameters changed.*/
+    if (lnode->param_source != NULL && lnode->base.needs_parameter_refresh)
+    {
+        /* pass NULL to forward: constant param_source never depends on u */
+        lnode->param_source->forward(lnode->param_source, NULL);
+    }
+
     refresh_param_values(lnode);
 
     expr *x = node->left;

src/atoms/affine/scalar_mult.c

@@ -28,7 +28,17 @@
 static void forward(expr *node, const double *u)
 {
     expr *child = node->left;
-    double a = ((scalar_mult_expr *) node)->param_source->value[0];
+    scalar_mult_expr *snode = (scalar_mult_expr *) node;
+
+    /* Refresh param_source expression tree if parameters changed.*/
+    if (node->needs_parameter_refresh)
+    {
+        /* pass NULL to forward: constant param_source never depends on u */
+        snode->param_source->forward(snode->param_source, NULL);
+        node->needs_parameter_refresh = false;
+    }
+
+    double a = snode->param_source->value[0];
 
     /* child's forward pass */
     child->forward(child, u);

src/atoms/affine/vector_mult.c

@@ -28,7 +28,17 @@
 static void forward(expr *node, const double *u)
 {
     expr *child = node->left;
-    const double *a = ((vector_mult_expr *) node)->param_source->value;
+    vector_mult_expr *vnode = (vector_mult_expr *) node;
+
+    /* Refresh param_source expression tree if parameters changed.*/
+    if (node->needs_parameter_refresh)
+    {
+        /* pass NULL to forward: constant param_source never depends on u */
+        vnode->param_source->forward(vnode->param_source, NULL);
+        node->needs_parameter_refresh = false;
+    }
+
+    const double *a = vnode->param_source->value;
 
     /* child's forward pass */
     child->forward(child, u);

tests/all_tests.c

@@ -49,6 +49,7 @@
 #include "jacobian_tests/other/test_prod_axis_zero.h"
 #include "jacobian_tests/other/test_quad_form.h"
 #include "numerical_diff/test_numerical_diff.h"
+#include "problem/test_param_broadcast.h"
 #include "problem/test_param_prob.h"
 #include "problem/test_problem.h"
 #include "utils/test_cblas.h"
@@ -358,6 +359,11 @@ int main(void)
     mu_run_test(test_param_right_matmul_rectangular, tests_run);
     mu_run_test(test_param_shared_left_matmul_problem, tests_run);
     mu_run_test(test_param_fixed_skip_in_update, tests_run);
+
+    printf("\n--- Parameter + Broadcast Tests ---\n");
+    mu_run_test(test_param_broadcast_vector_mult, tests_run);
+    mu_run_test(test_param_sum_scalar_mult, tests_run);
+    mu_run_test(test_param_broadcast_left_matmul, tests_run);
 #endif /* PROFILE_ONLY */
 
 #ifdef PROFILE_ONLY