[liuanji]

Summary

Three related problems with block-sparse sum edges, and with edges that list the same (sum block, child block) pair more than once:

1. block_sparse_rnd_blk_edge_constructor could connect a sum block to the same child block twice.
2. SumNodes.get_params(as_matrix = True) and get_param_flows(as_matrix = True) kept only one copy of such a repeated edge block, while the forward pass adds every copy.
3. SumNodes.set_params with a dense [num_nodes, num_ch_nodes] matrix ignored edge_ids, so it failed on any block-sparse sum. The 4-D dense form gave each copy of a repeated edge block its full value.

1. The random block-sparse constructor

The constructor drew child blocks with replacement: one random permutation for coverage, topped up with torch.randint draws, then split into groups per sum block. An HMM with 16 latents in blocks of 4 and num_chs_per_block = 2 got two repeated pairs among the eight edges of its transition sums:

edge_ids = [[0, 0, 1, 1, 2, 2, 3, 3],
            [3, 3, 1, 0, 3, 3, 0, 2]]     # (0, 3) and (2, 3) twice

Every sum block now gets distinct child blocks:

• Each child block is forced into one random slot, which keeps the documented guarantee that every child block is connected when there are enough slots.
• A sum block's other slots take its smallest random scores among the child blocks it doesn't have yet.
• Asking for more child blocks than exist connects to each of them once.

Over 3,000 seeds, every (sum block, child block) pair is chosen at the expected rate, within sampling noise.

2. The dense views

Both functions scattered the edge blocks with a plain assignment, which keeps one copy of a repeated pair. On the HMM above, the rows of the dense transition matrix summed to between 0.265 and 1, while the compiled circuit itself was normalized (log of the total probability over all strings: −2.8e-7). Both now accumulate the copies. For distinct edge blocks the result is unchanged.

3. set_params with a dense matrix

The [num_nodes, num_ch_nodes] form reshaped the matrix as if every (sum block, child block) pair were an edge. On a block-sparse sum, or one with a repeated edge block, it stored one parameter block per pair, and the normalization that follows failed (stack expects each tensor to be equal size). With normalize = False it silently kept parameters that don't match the edges. As a result, juice.structures.HMM(alpha = ...) failed with a block-sparse sum_edge_ids_constructor.

The [num_node_blocks, num_ch_node_blocks, block_size, ch_block_size] form did read the edge blocks, but gave every copy of a repeated pair its full value, so normalization counted it twice.

Both forms are now read at the edge blocks. An edge block listed k times gets 1/k of its value in every copy, so the copies add back up to it. A matrix set with set_params now comes back unchanged from get_params(as_matrix = True). Entries outside the edge blocks have no edge to go to, as before; the docstring now says so.

Compatibility

• The same seed now builds a different random block-sparse structure than before. Circuits saved with juice.save load unchanged, since their edges are stored explicitly. Code that builds a block-sparse PC once and saves it is unaffected.
• Fully connected sums store bit-identical parameters through set_params, which a test pins. Their edges are listed in the order the old reshape assumed. This is the path HMM(alpha = ...) / gamma take with dense edges.
• get_params(as_matrix = True) and get_param_flows(as_matrix = True) are unchanged when edge blocks are distinct.
• LVD's set_params call uses distinct edges, so its values are unchanged.
• Copy, juice.load and blockify pass per-edge parameters, a form this PR doesn't touch.

Tests

New file tests/nodes/block_sparse_edges_test.py, 22 tests:

• The constructor: never repeats a child block within a sum block, connects every child block when it can, and caps the request at the number of child blocks. It's checked on six configurations over 20 seeds each, including two child groups and fewer slots than child blocks.
• Dense parameters of a sum with a repeated edge block are normalized, and reproduce the compiled circuit's probabilities.
• Dense flows keep their total.
• The dense set_params round trip: both forms, with and without normalization, on a fully connected, a block-sparse and a repeated-edge sum.
• Unchanged storage for a fully connected sum.
• An HMM given alpha with block-sparse transitions: it reads alpha at the sampled edge blocks and compiles to a normalized circuit.

Six of the first commit's 8 tests fail on main. Seven of the second commit's 14 tests fail on the old set_params.

Full suite: 1275 passed, 26 skipped.

🤖 Generated with Claude Code