KimiK2.Config(
        channels_in=16,
        channels_out=32,
        num_layers=2,
        proj_in=Embedding.Config(
            channels_in=32,
            channels_out=16,
            padding_idx=None,
            device=None,
            dtype=None,
            lookup_dtype=None,
            shard='vocab',
            depth_index=(),
            init_weight=functools.partial(torch.nn.init.normal_, std=0.02)
        ),
        block=[
        │   TransformerBlock.Config(
        │       channels_in=16,
        │       channels_out=16,
        │       attn=MultiHeadLatentAttention.Config(
        │       │   channels_in=16,
        │       │   channels_out=16,
        │       │   num_heads=2,
        │       │   channels_qk_nope_head=4,
        │       │   channels_qk_rope_head=4,
        │       │   channels_v_head=4,
        │       │   q_lora_rank=None,
        │       │   kv_lora_rank=8,
        │       │   bias=False,
        │       │   dropout=0.0,
        │       │   causal=True,
        │       │   softmax_scale=None,
        │       │   rope=RoPE.Config(
        │       │       channels_head=4,
        │       │       reduction_mode='cat',
        │       │       frequencies=HuggingFaceFrequencies.Config(base=50000.0),
        │       │       dtype=torch.float32
        │       │   ),
        │       │   attn_kernel=LatentAttention.Config(
        │       │       absorb=True,
        │       │       attn_kernel=SdpaNaive.Config()
        │       │   ),
        │       │   norm_q_lora=RMSNorm.Config(
        │       │       channels_in=-1,
        │       │       channels_out=-1,
        │       │       eps=1e-06,
        │       │       elementwise_affine=True,
        │       │       device=None,
        │       │       dtype=None
        │       │   ),
        │       │   proj_q=Linear.Config(
        │       │       channels_in=16,
        │       │       channels_out=16,
        │       │       bias=False,
        │       │       depth_index=((0, 2),),
        │       │       device=None,
        │       │       dtype=None,
        │       │       init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │       │       init_bias=torch.nn.init.zeros_,
        │       │       linear_fn=<built-in function linear>,
        │       │       shard=None
        │       │   ),
        │       │   proj_q_a=Linear.Config(
        │       │       channels_in=-1,
        │       │       channels_out=-1,
        │       │       bias=False,
        │       │       depth_index=(),
        │       │       device=None,
        │       │       dtype=None,
        │       │       init_weight=priml.model.init.kaiming_uniform,
        │       │       init_bias=torch.nn.init.zeros_,
        │       │       linear_fn=<built-in function linear>,
        │       │       shard=None
        │       │   ),
        │       │   proj_q_b=Linear.Config(
        │       │       channels_in=-1,
        │       │       channels_out=-1,
        │       │       bias=False,
        │       │       depth_index=(),
        │       │       device=None,
        │       │       dtype=None,
        │       │       init_weight=priml.model.init.kaiming_uniform,
        │       │       init_bias=torch.nn.init.zeros_,
        │       │       linear_fn=<built-in function linear>,
        │       │       shard=None
        │       │   ),
        │       │   proj_kv_a=Linear.Config(
        │       │       channels_in=16,
        │       │       channels_out=12,
        │       │       bias=False,
        │       │       depth_index=((0, 2),),
        │       │       device=None,
        │       │       dtype=None,
        │       │       init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │       │       init_bias=torch.nn.init.zeros_,
        │       │       linear_fn=<built-in function linear>,
        │       │       shard=None
        │       │   ),
        │       │   proj_kv_b=Linear.Config(
        │       │       channels_in=8,
        │       │       channels_out=16,
        │       │       bias=False,
        │       │       depth_index=((0, 2),),
        │       │       device=None,
        │       │       dtype=None,
        │       │       init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │       │       init_bias=torch.nn.init.zeros_,
        │       │       linear_fn=<built-in function linear>,
        │       │       shard=None
        │       │   ),
        │       │   proj_out=Linear.Config(
        │       │       channels_in=8,
        │       │       channels_out=16,
        │       │       bias=False,
        │       │       depth_index=((0, 2),),
        │       │       device=None,
        │       │       dtype=None,
        │       │       init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │       │       init_bias=torch.nn.init.zeros_,
        │       │       linear_fn=<built-in function linear>,
        │       │       shard=None
        │       │   ),
        │       │   norm_kv_lora=RMSNorm.Config(
        │       │       channels_in=8,
        │       │       channels_out=8,
        │       │       eps=1e-06,
        │       │       elementwise_affine=True,
        │       │       device=None,
        │       │       dtype=None
        │       │   ),
        │       │   init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │       │   depth_index=((0, 2),),
        │       │   shard=None
        │       ),
        │       ffn=SwiGLU.Config(
        │           channels_in=16,
        │           channels_out=16,
        │           channels_hidden=32,
        │           expansion=2.6666666666666665,
        │           round_to=256,
        │           gate=True,
        │           bias=False,
        │           split_gate_projection=False,
        │           norm=None,
        │           act=priml.model.swiglu.silu,
        │           depth_index=((0, 2),),
        │           init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │           init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │           up_proj=None,
        │           middle=None,
        │           down_proj=None,
        │           shard='colwise'
        │       ),
        │       norm1=RMSNorm.Config(
        │           channels_in=16,
        │           channels_out=16,
        │           eps=1e-06,
        │           elementwise_affine=True,
        │           device=None,
        │           dtype=None
        │       ),
        │       norm2=RMSNorm.Config(
        │           channels_in=16,
        │           channels_out=16,
        │           eps=1e-06,
        │           elementwise_affine=True,
        │           device=None,
        │           dtype=None
        │       ),
        │       prenorm=True,
        │       checkpoint=False,
        │       depth_index=((0, 2),)
        │   ),
        │   TransformerBlock.Config(
        │       channels_in=16,
        │       channels_out=16,
        │       attn=MultiHeadLatentAttention.Config(
        │       │   channels_in=16,
        │       │   channels_out=16,
        │       │   num_heads=2,
        │       │   channels_qk_nope_head=4,
        │       │   channels_qk_rope_head=4,
        │       │   channels_v_head=4,
        │       │   q_lora_rank=None,
        │       │   kv_lora_rank=8,
        │       │   bias=False,
        │       │   dropout=0.0,
        │       │   causal=True,
        │       │   softmax_scale=None,
        │       │   rope=RoPE.Config(
        │       │       channels_head=4,
        │       │       reduction_mode='cat',
        │       │       frequencies=HuggingFaceFrequencies.Config(base=50000.0),
        │       │       dtype=torch.float32
        │       │   ),
        │       │   attn_kernel=LatentAttention.Config(
        │       │       absorb=True,
        │       │       attn_kernel=SdpaNaive.Config()
        │       │   ),
        │       │   norm_q_lora=RMSNorm.Config(
        │       │       channels_in=-1,
        │       │       channels_out=-1,
        │       │       eps=1e-06,
        │       │       elementwise_affine=True,
        │       │       device=None,
        │       │       dtype=None
        │       │   ),
        │       │   proj_q=Linear.Config(
        │       │       channels_in=16,
        │       │       channels_out=16,
        │       │       bias=False,
        │       │       depth_index=((1, 2),),
        │       │       device=None,
        │       │       dtype=None,
        │       │       init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │       │       init_bias=torch.nn.init.zeros_,
        │       │       linear_fn=<built-in function linear>,
        │       │       shard=None
        │       │   ),
        │       │   proj_q_a=Linear.Config(
        │       │       channels_in=-1,
        │       │       channels_out=-1,
        │       │       bias=False,
        │       │       depth_index=(),
        │       │       device=None,
        │       │       dtype=None,
        │       │       init_weight=priml.model.init.kaiming_uniform,
        │       │       init_bias=torch.nn.init.zeros_,
        │       │       linear_fn=<built-in function linear>,
        │       │       shard=None
        │       │   ),
        │       │   proj_q_b=Linear.Config(
        │       │       channels_in=-1,
        │       │       channels_out=-1,
        │       │       bias=False,
        │       │       depth_index=(),
        │       │       device=None,
        │       │       dtype=None,
        │       │       init_weight=priml.model.init.kaiming_uniform,
        │       │       init_bias=torch.nn.init.zeros_,
        │       │       linear_fn=<built-in function linear>,
        │       │       shard=None
        │       │   ),
        │       │   proj_kv_a=Linear.Config(
        │       │       channels_in=16,
        │       │       channels_out=12,
        │       │       bias=False,
        │       │       depth_index=((1, 2),),
        │       │       device=None,
        │       │       dtype=None,
        │       │       init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │       │       init_bias=torch.nn.init.zeros_,
        │       │       linear_fn=<built-in function linear>,
        │       │       shard=None
        │       │   ),
        │       │   proj_kv_b=Linear.Config(
        │       │       channels_in=8,
        │       │       channels_out=16,
        │       │       bias=False,
        │       │       depth_index=((1, 2),),
        │       │       device=None,
        │       │       dtype=None,
        │       │       init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │       │       init_bias=torch.nn.init.zeros_,
        │       │       linear_fn=<built-in function linear>,
        │       │       shard=None
        │       │   ),
        │       │   proj_out=Linear.Config(
        │       │       channels_in=8,
        │       │       channels_out=16,
        │       │       bias=False,
        │       │       depth_index=((1, 2),),
        │       │       device=None,
        │       │       dtype=None,
        │       │       init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │       │       init_bias=torch.nn.init.zeros_,
        │       │       linear_fn=<built-in function linear>,
        │       │       shard=None
        │       │   ),
        │       │   norm_kv_lora=RMSNorm.Config(
        │       │       channels_in=8,
        │       │       channels_out=8,
        │       │       eps=1e-06,
        │       │       elementwise_affine=True,
        │       │       device=None,
        │       │       dtype=None
        │       │   ),
        │       │   init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │       │   depth_index=((1, 2),),
        │       │   shard=None
        │       ),
        │       ffn=MoE.Config(
        │       │   channels_in=16,
        │       │   channels_out=16,
        │       │   router=SigmoidRouter.Config(
        │       │       channels_in=16,
        │       │       num_experts=2,
        │       │       top_k=1,
        │       │       norm_topk_prob=True,
        │       │       use_correction_bias=True,
        │       │       routed_scaling_factor=2.0,
        │       │       n_group=1,
        │       │       topk_group=1
        │       │   ),
        │       │   expert=SwiGLU.Config(
        │       │       channels_in=16,
        │       │       channels_out=16,
        │       │       channels_hidden=16,
        │       │       expansion=2.6666666666666665,
        │       │       round_to=256,
        │       │       gate=True,
        │       │       bias=False,
        │       │       split_gate_projection=False,
        │       │       norm=None,
        │       │       act=priml.model.swiglu.silu,
        │       │       depth_index=((1, 2),),
        │       │       init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │       │       init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │       │       up_proj=None,
        │       │       middle=None,
        │       │       down_proj=None,
        │       │       shard='colwise'
        │       │   ),
        │       │   num_shared_experts=1,
        │       │   shared_expert=SwiGLU.Config(
        │       │       channels_in=16,
        │       │       channels_out=16,
        │       │       channels_hidden=16,
        │       │       expansion=2.6666666666666665,
        │       │       round_to=256,
        │       │       gate=True,
        │       │       bias=False,
        │       │       split_gate_projection=False,
        │       │       norm=None,
        │       │       act=priml.model.swiglu.silu,
        │       │       depth_index=((1, 2),),
        │       │       init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
        │       │       init_weight_out=functools.partial(torch.nn.init.normal_, std=0.02),
        │       │       up_proj=None,
        │       │       middle=None,
        │       │       down_proj=None,
        │       │       shard='colwise'
        │       │   ),
        │       │   aux_loss_weight=0.01,
        │       │   depth_index=((1, 2),)
        │       ),
        │       norm1=RMSNorm.Config(
        │           channels_in=16,
        │           channels_out=16,
        │           eps=1e-06,
        │           elementwise_affine=True,
        │           device=None,
        │           dtype=None
        │       ),
        │       norm2=RMSNorm.Config(
        │           channels_in=16,
        │           channels_out=16,
        │           eps=1e-06,
        │           elementwise_affine=True,
        │           device=None,
        │           dtype=None
        │       ),
        │       prenorm=True,
        │       checkpoint=False,
        │       depth_index=((1, 2),)
        │   )
        ],
        proj_out=Sequential.Config(
            channels_in=16,
            channels_out=32,
            elements=[
                RMSNorm.Config(
                    channels_in=16,
                    channels_out=16,
                    eps=1e-06,
                    elementwise_affine=True,
                    device=None,
                    dtype=None
                ),
                Linear.Config(
                    channels_in=16,
                    channels_out=32,
                    bias=False,
                    depth_index=(),
                    device=None,
                    dtype=None,
                    init_weight=functools.partial(torch.nn.init.normal_, std=0.02),
                    init_bias=torch.nn.init.zeros_,
                    linear_fn=<built-in function linear>,
                    shard='vocab'
                )
            ],
            repeat=1,
            depth_index=()
        ),
        channels_hidden_dense=32,
        channels_hidden_expert=16,
        first_k_dense_replace=1
    )
