Repository navigation
feat(nodes)!: BlockScaleSumParams applies the z-correction (dual-flow EM) by default - #63
Merged
Merged
Conversation
… EM) by default apply_z_correction now defaults to True: a gated sum node gets the exact M-step under its gate unless it passes apply_z_correction = False. The backward accumulates the denominator flow F- (in gate space, ch_block_size times smaller than param_flows; about 1% of an EM step) and FullBatchEM / MiniBatchEM / Anemone update theta <- normalize(theta * F+ / F-). With False the layer allocates nothing and takes the plain M-step, as before. Breaking: a gate finer than ns.block_size along the node axis is not served by the correction, and its backward raises NotImplementedError, now naming apply_z_correction = False as the way out. Gradient optimizers still warn on a circuit that requests F-, so a gated model trained with SGD / Adam now warns unless built with False. - Tests: the default asserted on; the node-axis gate flow test passes False where the gate is finer than the node block, as does its CUDA-graph capture test; the mixed plain / gated sum group now runs with F- allocated.
liuanji
added a commit
that referenced
this pull request
Oct 22, 2026
Brings in #62: the sum layers' precision option, the cuBLAS dense path, the tf32x3 backward, and the shared choice of the tlmm fast path. #63, the z-correction default, is on this branch already as 1ed9d48. Conflicts: tensorcircuit.py and block_scale.py, both split into modules on this branch. Both take the branch's side. tensorcircuit.py adds #62's import of PRECISIONS and _pop_removed_precision_flags, and #62's change to the element-flow helper goes to model/inspection.py, where that helper now lives. block_scale.py's split modules already carry #63.
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
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.
BlockScaleSumParams(apply_z_correction = ...)now defaults toTrue: a gated sum node gets the exact M-step under its gate unless it opts out.What the default now does. The backward also accumulates the denominator flow
F-, andFullBatchEM/MiniBatchEM/Anemoneupdatetheta <- normalize(theta * F+ / F-).F-is kept in gate space,ch_block_sizetimes smaller thanparam_flows.With
apply_z_correction = Falsethe layer allocates nothing and takes the plain M-steptheta <- normalize(F+), as before.Breaking
block_size < ns.block_size) are not served by the correction. Their backward raisesNotImplementedError, whose message now says to passapply_z_correction = False.F-. A gated model trained with them now warns unless it is built withFalse.Tests
denom_param_flows_test: the default is asserted on, and the mixed plain / gated sum group now runs withF-allocated.test_a_node_axis_gate_computes_the_right_flowspassesFalsewhere the gate is finer than the node block, as doestest_a_node_axis_gate_is_captured.tests/optim, 773 passed, 12 skipped. The three edited files were re-run after the fix. None of them is marked slow.