LatentRecursion.Config(
        vocab_size=5,
        puzzle_grid_shape=(2, 3),
        channels_in=12,
        num_layers=1,
        num_heads=3,
        slow_cycles=2,
        fast_cycles=2,
        q_head_outputs=2,
        num_puzzle_identifiers=2,
        puzzle_emb_len=2,
        puzzle_emb_ndim=12,
        puzzle_emb_batch_size=2,
        puzzle_emb_init_std=0.0,
        pos2d_grid_shape=(2, 3),
        pos2d_box_shape=(2, 1),
        pos2d_init_std=1.0,
        feedback_init_std=0.0,
        compile=False,
        dtype=None,
        block=TransformerBlock.Config(
        │   channels_in=12,
        │   channels_out=12,
        │   attn=Attention.Config(
        │       channels_in=12,
        │       channels_out=12,
        │       num_heads=3,
        │       channels_head=4,
        │       num_heads_kv=3,
        │       bias=False,
        │       dropout=0.0,
        │       causal=False,
        │       rope=None,
        │       rotate=None,
        │       norm_qk=RMSNorm.Config(
        │           channels_in=4,
        │           channels_out=4,
        │           eps=1e-06,
        │           elementwise_affine=False,
        │           device=None,
        │           dtype=None
        │       ),
        │       share_qk_norm=True,
        │       norm_out=None,
        │       split_qkv_projection=False,
        │       proj_qkv=None,
        │       proj_out=None,
        │       init_weight=priml.baselines.sudoku.latent_recursion.truncated_normal_corrected,
        │       depth_index=(),
        │       attn_kernel=SdpaFused.Config()
        │   ),
        │   ffn=SwiGLU.Config(
        │       channels_in=12,
        │       channels_out=12,
        │       channels_hidden=256,
        │       expansion=2.6666666666666665,
        │       round_to=256,
        │       gate=True,
        │       bias=False,
        │       split_gate_projection=False,
        │       norm=RMSNorm.Config(
        │           channels_in=256,
        │           channels_out=256,
        │           eps=1e-06,
        │           elementwise_affine=False,
        │           device=None,
        │           dtype=None
        │       ),
        │       act=priml.model.swiglu.silu,
        │       depth_index=(),
        │       init_weight=priml.baselines.sudoku.latent_recursion.truncated_normal_corrected,
        │       init_weight_out=priml.baselines.sudoku.latent_recursion.truncated_normal_corrected,
        │       up_proj=None,
        │       middle=None,
        │       down_proj=None,
        │       shard='colwise'
        │   ),
        │   norm1=RMSNorm.Config(
        │       channels_in=12,
        │       channels_out=12,
        │       eps=1e-05,
        │       elementwise_affine=False,
        │       device=None,
        │       dtype=None
        │   ),
        │   norm2=RMSNorm.Config(
        │       channels_in=12,
        │       channels_out=12,
        │       eps=1e-05,
        │       elementwise_affine=False,
        │       device=None,
        │       dtype=None
        │   ),
        │   prenorm=False,
        │   checkpoint=False,
        │   depth_index=()
        )
    )
