load_reference_state_dict

load_reference_state_dict(mixer, state_dict)

Load a reference TriHyena checkpoint into QKVSequenceMixer(TriangleHyena).

The original TriHyena TriangleHyena module stores linear_x12i (x1 and x2 stacked), linear_vj, a (3D, 3) causal short-conv weight over [x1 | x2 | v], siren_kernel.*, to_out and a pre-norm. This maps them onto qkv_proj ([x1 | v | x2] channel blocks), short_conv_weight (same remap, transposed to [3, 3D]), global_conv.kernel and out_proj, and zeroes qkv_proj.bias and global_conv.shortcut when the mixer has them (the reference has neither; build with qkv_bias=False and use_shortcut=False to drop them).

The mixer must be built to match: SIRENKernelND(L_cache=2 * L - 1) with the reference L (its SIREN grid keeps the reference tap spacing 1 / (2 (L - 1)) at every runtime length and spans [-0.5, 0.5] only at L), mask_cfg = BoxModulationND(max_offset=<reference max_offset>) or Identity when it was 0, short_conv_kernel_size=3 and out_proj_bias=True.

Parameters:

• mixer (QKVSequenceMixer) – QKVSequenceMixer whose mixer is a TriangleHyena.

• state_dict (Mapping[str, Tensor]) – Reference state dict, optionally prefixed with core..

Returns:

The reference pre-norm tensors (weight, bias) for the caller to load into the residual block’s LayerNorm.

Raises:

• ValueError – If a tensor’s shape does not match the mixer’s, or a reference key is missing.

• ValueError – If a reference key other than grid_cache is unused.

• ValueError – If a mixer parameter has no reference counterpart.

Return type:

dict[str, Tensor]