TransformerBlock.Config(
        channels_in=16,
        channels_out=16,
        attn=Attention.Config(
            channels_in=16,
            channels_out=16,
            num_heads=2,
            channels_head=8,
            num_heads_kv=2,
            bias=False,
            dropout=0.0,
            causal=False,
            rope=None,
            rotate=None,
            norm_qk=None,
            share_qk_norm=True,
            norm_out=None,
            split_qkv_projection=False,
            proj_qkv=None,
            proj_out=None,
            init_weight=priml.model.init.kaiming_uniform,
            depth_index=(),
            attn_kernel=SdpaFused.Config()
        ),
        ffn=SwiGLU.Config(
            channels_in=16,
            channels_out=16,
            channels_hidden=256,
            expansion=2.6666666666666665,
            round_to=256,
            gate=True,
            bias=False,
            split_gate_projection=False,
            norm=None,
            act=priml.model.swiglu.silu,
            depth_index=(),
            init_weight=priml.model.init.unit_fan_in_uniform,
            init_weight_out=torch.nn.init.zeros_,
            up_proj=None,
            middle=None,
            down_proj=None,
            shard='colwise'
        ),
        norm1=RMSNorm.Config(
            channels_in=16,
            channels_out=16,
            eps=1e-06,
            elementwise_affine=False,
            device=None,
            dtype=None
        ),
        norm2=RMSNorm.Config(
            channels_in=16,
            channels_out=16,
            eps=1e-06,
            elementwise_affine=False,
            device=None,
            dtype=None
        ),
        prenorm=True,
        checkpoint=False,
        depth_index=()
    )
