Skip to content

Refresh composite param_source subtrees on parameter updates - #107

Merged
dance858 merged 1 commit into
mainfrom
param-source-mark-refresh
Jul 13, 2026
Merged

dance858 merged 1 commit into
mainfrom
param-source-mark-refresh

Conversation

@Transurgeon

Copy link
Copy Markdown
Member

Problem

A param_source subtree is evaluated by the owning atom's gated recursion, not the main forward walk, so expr_set_needs_refresh never reaches nodes inside a composite source. Gated nodes there — promote, nested mults, cached matmul coefficients — keep serving values cached at construction after problem_update_params.

The failing shapes are common, not exotic (each verified with a standalone repro):

lowered shape example cvxpy origin result before this PR
vector_mult(promote(p), A) feeding left_matmul (p * A) @ x stale A on every re-solve
vector_mult(promote(g), Sig) feeding quad_form quad_form(x, g * Sig) stale Q
division/scaling chains feeding quad_form quad_over_lin(x, p) (canonicalized to I/p) stale Q
two nested gated levels ((2*p) * A) @ x stale A

Shapes that were already correct and stay untouched: bare-Parameter sources (refreshed by problem_update_params' memcpy), purely additive composites (p1+p2 — ungated nodes recurse unconditionally), and computed multipliers of variables ((2*p)*x — the gated node sits in the main tree, which the walker reaches).

Fix

One line per source-owning atom (scalar_mult, vector_mult, left_matmul — which also serves right_matmul — convolve, kron, quad_form): mark the side subtree dirty at the moment the atom is about to re-evaluate it,

if (node->param_source != NULL && node->needs_parameter_refresh)
{
    expr_set_needs_refresh(node->param_source);   /* new */
    node->param_source->forward(node->param_source, NULL);
}

The owner's flag is already set by the existing walker (it lives in the main tree), so the mark happens exactly once per parameter update; nested gated levels recurse for free because each level applies the same line to its own source. No struct changes, no walker changes, no behavior change for one-shot solves or bare-parameter sources.

Tests

tests/problem/test_param_source_refresh.h pins the three failing shapes (composite matmul coefficient, composite quad matrix, two nested gated levels), asserting constraint values and Jacobian/gradient values across an update. Each fails without the fix (e.g. stale 10.0 where 50.0 is expected) and passes with it. Full suite: 417/417.

Downstream validation (cvxpy pr-a/b/c stack, cached re-solve paths): the cvxpy-side canaries (quad_over_lin symbolic quad matrix, scaled parameter coefficient) and all four repro shapes above go from stale to fresh on this branch with no other changes; full cvxpy test suite green (modulo pre-existing, unrelated failures verified identical on other builds).

Relation to #104

This supersedes the first commit of #104 (param-source refresh propagation) with a smaller, atom-local change: no param_source pointer hoisted into the base expr struct. #104's second commit (retaining registered parameter nodes) is deliberately not included — it is an independent API-soundness hardening that can be its own PR if wanted.

🤖 Generated with Claude Code

@Transurgeon
Transurgeon force-pushed the param-source-mark-refresh branch from c93a83c to 4248f7d Compare July 12, 2026 22:16
…them

A param_source subtree is evaluated by the owning atom's gated recursion,
not the main forward walk, so expr_set_needs_refresh never reaches nodes
inside a composite source. Gated nodes there (promote, nested mults,
cached matmul coefficients) kept serving values cached at build time
after problem_update_params: p*A @ x and quad_form(x, g*Sig) both solved
with the original parameter values on every re-solve.

Fix: each source-owning atom marks its subtree dirty at the moment it is
about to re-evaluate it (one expr_set_needs_refresh call per atom, gated
on the flag the main walker already sets). Nested gated levels recurse
for free. No struct or walker changes; no behavior change for one-shot
solves or bare-parameter sources.

Tests: tests/problem/test_param_source_refresh.h pins the three failing
shapes (composite matmul coefficient, composite quad matrix, two nested
gated levels); each fails without the fix and passes with it.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@dance858

Copy link
Copy Markdown
Collaborator

Beautiful solution. I could stare at it for hours.

@dance858
dance858 merged commit 4172c5e into main Jul 13, 2026
11 checks passed
@Transurgeon
Transurgeon deleted the param-source-mark-refresh branch July 13, 2026 03:30
Transurgeon added a commit that referenced this pull request Jul 13, 2026
…them (#107)

A param_source subtree is evaluated by the owning atom's gated recursion,
not the main forward walk, so expr_set_needs_refresh never reaches nodes
inside a composite source. Gated nodes there (promote, nested mults,
cached matmul coefficients) kept serving values cached at build time
after problem_update_params: p*A @ x and quad_form(x, g*Sig) both solved
with the original parameter values on every re-solve.

Fix: each source-owning atom marks its subtree dirty at the moment it is
about to re-evaluate it (one expr_set_needs_refresh call per atom, gated
on the flag the main walker already sets). Nested gated levels recurse
for free. No struct or walker changes; no behavior change for one-shot
solves or bare-parameter sources.

Tests: tests/problem/test_param_source_refresh.h pins the three failing
shapes (composite matmul coefficient, composite quad matrix, two nested
gated levels); each fails without the fix and passes with it.

Co-authored-by: Transurgeon <peter.zijie@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
dance858 added a commit that referenced this pull request Sep 19, 2026
… walks it

Since the coefficient atoms install set_needs_refresh_children to reach
param_source, the problem-level walk re-arms the whole source subtree
before forward runs. The expr_set_needs_refresh(param_source) call inside
the gated forward branch (from #107, when the walk could not see
param_source) walked the same nodes a second time and found nothing left
to clear. Remove it from scalar_mult, vector_mult, kron, convolve,
left_matmul and quad_form; keep the gate, the source forward and the
flag reset.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Transurgeon added a commit that referenced this pull request Oct 4, 2026
…n hook (#125)

* Prune parameter-free subtrees; reach param_source through the children hook

problem_update_params floods expr_set_needs_refresh over the whole DAG,
clearing work->jacobian_evaluated on every node. That re-arms the affine
bump-skip in eval_jacobian even for subtrees no parameter can reach, so
their Jacobians are recomputed after every parameter update despite being
provably unchanged.

Give expr a memoized has_params and let expr_set_needs_refresh compute it
while it walks, returning it to the caller. It starts true so nothing is
pruned before the first walk; that first walk is the one that discovers the
dependency and behaves exactly as before. Every walk after it skips
parameter-free subtrees outright. This also stops the unguarded recursion
from re-walking shared parameter-free subtrees once per path.

The answer is computed, not declared. param_source is a child living
outside left/right -- structurally the same as hstack's args[] -- so the
existing set_needs_refresh_children hook is extended to cover it. The six
coefficient atoms (convolve, kron, left_matmul, scalar_mult, vector_mult,
quad_form) walk their param_source through that hook, and a parameter leaf
reports its own param_id >= 0. right_matmul needs nothing; its
parameterized path delegates to new_left_matmul_dense.

Computing rather than declaring matters here: param_source holds whichever
operand is variable-free, which is often a plain PARAM_FIXED constant (see
test_constant_broadcast_vector_mult and friends). Seeding those atoms
"contains a parameter" would be pessimistic; walking the source gets it
right, so vector_mult(const, x) now prunes while vector_mult(param, x)
still re-marks.

Dense lasso, m=2000 n=785, 50-point lambda path (tests/profiling/
profile_lasso.h, A constant, lambda the only registered parameter):

  gradient   4.64 -> 1.34 ms/iter   (3.5x; Ax Jacobian is 1.57M nnz)
  jacobian   0.007 -> 0.001 ms/iter
  forward    0.260 -> 0.249 ms/iter (unchanged, as expected)

Objective values are bit-identical across the sweep. sizeof(expr) is
unchanged at 192 bytes -- the flag packs into existing padding.

Tests. Every mutation of this change is now pinned by a test that fails
without it, verified by applying each mutation against a clean build:

  delete the prune guard          -> test_refresh_prunes_param_free
  init has_params = false         -> test_values_version_affine
  memoize true instead of child   -> test_refresh_prunes_param_free
  parameter hook ignores param_id -> test_refresh_prunes_fixed_constant
  hstack hook returns false       -> test_values_version_param_under_hstack
  drop convolve's hook            -> test_param_scalar_mult_convolve
  drop kron's hook                -> test_composite_source_kron
  drop left_matmul's hook         -> test_composite_source_left_matmul
  drop quad_form's hook           -> test_composite_source_quad_form
  drop scalar_mult's hook         -> test_values_version_param_under_hstack
  drop vector_mult's hook         -> test_composite_source_left_matmul

Six of those were undetected before this commit's test changes. The reason
is structural: parameter-dependence is memoized on the FIRST refresh walk,
so a subtree the walk fails to reach is still marked correctly that once and
only goes stale from the SECOND update onward. Almost every parameter test
called problem_update_params exactly once, so the suite could not see it.
The fix is to give the composite param_source tests a second update --
restoring the original parameter value, so they assert against numbers
already in the test -- rather than to add parallel tests.

Four new tests: a parameter-free node keeps its latch armed across repeated
walks (proven by poisoning its Jacobian and observing the poison survives);
a parameter-dependent node is re-armed by every walk; and a matched pair
showing a PARAM_FIXED constant prunes while an updatable parameter of the
same shape does not.

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

* Drop the forward-time param_source re-mark; the children hook already walks it

Since the coefficient atoms install set_needs_refresh_children to reach
param_source, the problem-level walk re-arms the whole source subtree
before forward runs. The expr_set_needs_refresh(param_source) call inside
the gated forward branch (from #107, when the walk could not see
param_source) walked the same nodes a second time and found nothing left
to clear. Remove it from scalar_mult, vector_mult, kron, convolve,
left_matmul and quad_form; keep the gate, the source forward and the
flag reset.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* Document the set_needs_refresh_children contract; fix a test comment

The hook comment in expr.h still described the pre-pruning design
(hstack only, NULL for unary/binary atoms) and never said that the hook
must report whether the nodes it reaches hold an updatable parameter.
Since the forward-time re-mark is gone, the hook is the only way the
refresh walk reaches a param_source, and its return is what keeps the
owner's has_params true, so spell that out. Also correct the
needs_parameter_refresh comment: the flag is set only on nodes whose
subtree contains an updatable parameter, not on all nodes.

test_composite_source_nested_gates sets theta[0] = 1.0 on its last
round; the comment said "p back to 2".

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: dance858 <danielcederberg1@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants