Prune parameter-free subtrees; reach param_source through the children hook - #125
Open
Transurgeon wants to merge 3 commits into
Open
Transurgeon wants to merge 3 commits into
Transurgeon wants to merge 3 commits into
Conversation
Transurgeon
force-pushed
the
prune-param-free-refresh
branch
from
September 19, 2026 14:59
14388e8 to
4c18a31
Compare
Transurgeon
changed the title
…n 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>
Transurgeon
force-pushed
the
prune-param-free-refresh
branch
from
September 19, 2026 15:23
4c18a31 to
11c3cb8
Compare
… 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>
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>
dance858
self-requested a review
September 19, 2026 20:22
Collaborator
|
Looks great. Very simple solution :) |
dance858
approved these changes
Sep 19, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.
No existing test needed changing; two new ones pin the behaviour -- a parameter-free node keeps its latch armed across repeated walks (proven by poisoning its Jacobian values and observing they survive), and a parameter-dependent node is re-armed by every walk and tracks a real parameter change.