Family.Config(
        name='0719_V1',
        training=WellboreFoldLoop.Config(
        │   study_name='wellbore_geology',
        │   experiment_name='exp000',
        │   base_dir=PosixPath('/opt/scratch'),
        │   working_dir=PosixPath('/opt/scratch/opt/scratch/runs/wellbore_geology/exp000/0719_V1'),
        │   doc='',
        │   step=WellboreTrainStep.Config(
        │   │   model=SequenceModel.Config(
        │   │   │   channels_in=16,
        │   │   │   channels_hidden=32,
        │   │   │   kernel_size=-1,
        │   │   │   dropout=0.05,
        │   │   │   reference_columns=400,
        │   │   │   reference_window=100.0,
        │   │   │   target_stats=(1.591691, 15.853182),
        │   │   │   channels_last=True,
        │   │   │   unet=PosteriorUNet.Config(
        │   │   │   │   channels_in=16,
        │   │   │   │   channels_out=32,
        │   │   │   │   channels_hidden=(96, 192, 384, 768),
        │   │   │   │   stem=ResidualBlock.Config(
        │   │   │   │   │   channels_in=16,
        │   │   │   │   │   channels_out=32,
        │   │   │   │   │   conv1=ConvNormAct.Config(
        │   │   │   │   │       channels_in=16,
        │   │   │   │   │       channels_out=32,
        │   │   │   │   │       conv=Conv2d.Config(
        │   │   │   │   │           channels_in=16,
        │   │   │   │   │           channels_out=32,
        │   │   │   │   │           kernel_size=3,
        │   │   │   │   │           stride=1,
        │   │   │   │   │           padding='same',
        │   │   │   │   │           dilation=1,
        │   │   │   │   │           groups=1,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           depth_index=(),
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None,
        │   │   │   │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │           init_bias=torch.nn.init.zeros_
        │   │   │   │   │       ),
        │   │   │   │   │       norm=BatchNorm2d.Config(
        │   │   │   │   │           channels_in=32,
        │   │   │   │   │           channels_out=32,
        │   │   │   │   │           momentum=0.1,
        │   │   │   │   │           eps=1e-05,
        │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None
        │   │   │   │   │       ),
        │   │   │   │   │       activation=priml.model.swiglu.silu,
        │   │   │   │   │       drop=Dropout2d.Config(
        │   │   │   │   │           channels_in=-1,
        │   │   │   │   │           channels_out=-1,
        │   │   │   │   │           p=0.05,
        │   │   │   │   │           inplace=False
        │   │   │   │   │       )
        │   │   │   │   │   ),
        │   │   │   │   │   conv2=Conv2d.Config(
        │   │   │   │   │       channels_in=32,
        │   │   │   │   │       channels_out=32,
        │   │   │   │   │       kernel_size=3,
        │   │   │   │   │       stride=1,
        │   │   │   │   │       padding='same',
        │   │   │   │   │       dilation=1,
        │   │   │   │   │       groups=1,
        │   │   │   │   │       bias=False,
        │   │   │   │   │       depth_index=(),
        │   │   │   │   │       device=None,
        │   │   │   │   │       dtype=None,
        │   │   │   │   │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │       init_bias=torch.nn.init.zeros_
        │   │   │   │   │   ),
        │   │   │   │   │   norm=BatchNorm2d.Config(
        │   │   │   │   │       channels_in=32,
        │   │   │   │   │       channels_out=32,
        │   │   │   │   │       momentum=0.1,
        │   │   │   │   │       eps=1e-05,
        │   │   │   │   │       elementwise_affine=True,
        │   │   │   │   │       device=None,
        │   │   │   │   │       dtype=None
        │   │   │   │   │   ),
        │   │   │   │   │   drop=Dropout2d.Config(
        │   │   │   │   │       channels_in=-1,
        │   │   │   │   │       channels_out=-1,
        │   │   │   │   │       p=0.05,
        │   │   │   │   │       inplace=False
        │   │   │   │   │   ),
        │   │   │   │   │   activation=priml.model.swiglu.silu,
        │   │   │   │   │   skip=Conv2d.Config(
        │   │   │   │   │       channels_in=16,
        │   │   │   │   │       channels_out=32,
        │   │   │   │   │       kernel_size=1,
        │   │   │   │   │       stride=1,
        │   │   │   │   │       padding='same',
        │   │   │   │   │       dilation=1,
        │   │   │   │   │       groups=1,
        │   │   │   │   │       bias=False,
        │   │   │   │   │       depth_index=(),
        │   │   │   │   │       device=None,
        │   │   │   │   │       dtype=None,
        │   │   │   │   │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │       init_bias=torch.nn.init.zeros_
        │   │   │   │   │   )
        │   │   │   │   ),
        │   │   │   │   stem_down=AvgPool2d.Config(
        │   │   │   │       channels_in=-1,
        │   │   │   │       channels_out=-1,
        │   │   │   │       kernel_size=(2, 4),
        │   │   │   │       stride=(2, 4),
        │   │   │   │       padding=0,
        │   │   │   │       ceil_mode=True,
        │   │   │   │       count_include_pad=True,
        │   │   │   │       divisor_override=None
        │   │   │   │   ),
        │   │   │   │   to_backbone=Conv2d.Config(
        │   │   │   │       channels_in=32,
        │   │   │   │       channels_out=96,
        │   │   │   │       kernel_size=1,
        │   │   │   │       stride=1,
        │   │   │   │       padding='same',
        │   │   │   │       dilation=1,
        │   │   │   │       groups=1,
        │   │   │   │       bias=True,
        │   │   │   │       depth_index=(),
        │   │   │   │       device=None,
        │   │   │   │       dtype=None,
        │   │   │   │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       init_bias=None
        │   │   │   │   ),
        │   │   │   │   backbone=Backbone.Config(
        │   │   │   │   │   network=ConvNeXt.Config(
        │   │   │   │   │   │   channels_in=3,
        │   │   │   │   │   │   channels_hidden=(96, 192, 384, 768),
        │   │   │   │   │   │   blocks_per_stage=(3, 3, 27, 3),
        │   │   │   │   │   │   stem=ConvNeXtStem.Config(
        │   │   │   │   │   │       channels_in=3,
        │   │   │   │   │   │       channels_out=96,
        │   │   │   │   │   │       conv=Conv2d.Config(
        │   │   │   │   │   │           channels_in=3,
        │   │   │   │   │   │           channels_out=96,
        │   │   │   │   │   │           kernel_size=4,
        │   │   │   │   │   │           stride=4,
        │   │   │   │   │   │           padding=0,
        │   │   │   │   │   │           dilation=1,
        │   │   │   │   │   │           groups=1,
        │   │   │   │   │   │           bias=True,
        │   │   │   │   │   │           depth_index=(),
        │   │   │   │   │   │           device=None,
        │   │   │   │   │   │           dtype=None,
        │   │   │   │   │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │           init_bias=None
        │   │   │   │   │   │       ),
        │   │   │   │   │   │       norm=LayerNorm2d.Config(
        │   │   │   │   │   │           channels_in=96,
        │   │   │   │   │   │           channels_out=96,
        │   │   │   │   │   │           eps=1e-06,
        │   │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │   │           device=None,
        │   │   │   │   │   │           dtype=None
        │   │   │   │   │   │       )
        │   │   │   │   │   │   ),
        │   │   │   │   │   │   stage=[
        │   │   │   │   │   │   │   ConvNeXtStage.Config(
        │   │   │   │   │   │   │       channels_in=96,
        │   │   │   │   │   │   │       channels_out=96,
        │   │   │   │   │   │   │       blocks_per_stage=2,
        │   │   │   │   │   │   │       downsample=Identity.Config(
        │   │   │   │   │   │   │           channels_in=96,
        │   │   │   │   │   │   │           channels_out=96
        │   │   │   │   │   │   │       ),
        │   │   │   │   │   │   │       block=[
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=96,
        │   │   │   │   │   │   │       │       channels_out=96,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=96,
        │   │   │   │   │   │   │       │           channels_out=96,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=96,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=96,
        │   │   │   │   │   │   │       │           channels_out=96,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=96,
        │   │   │   │   │   │   │       │           channels_out=96,
        │   │   │   │   │   │   │       │           channels_hidden=384,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=96,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=96,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=96,
        │   │   │   │   │   │   │       │       channels_out=96,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=96,
        │   │   │   │   │   │   │       │           channels_out=96,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=96,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=96,
        │   │   │   │   │   │   │       │           channels_out=96,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=96,
        │   │   │   │   │   │   │       │           channels_out=96,
        │   │   │   │   │   │   │       │           channels_hidden=384,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=96,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=96,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=96,
        │   │   │   │   │   │   │       │       channels_out=96,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=96,
        │   │   │   │   │   │   │       │           channels_out=96,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=96,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=96,
        │   │   │   │   │   │   │       │           channels_out=96,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=96,
        │   │   │   │   │   │   │       │           channels_out=96,
        │   │   │   │   │   │   │       │           channels_hidden=384,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=96,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=96,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   )
        │   │   │   │   │   │   │       ]
        │   │   │   │   │   │   │   ),
        │   │   │   │   │   │   │   ConvNeXtStage.Config(
        │   │   │   │   │   │   │       channels_in=96,
        │   │   │   │   │   │   │       channels_out=192,
        │   │   │   │   │   │   │       blocks_per_stage=2,
        │   │   │   │   │   │   │       downsample=Sequential.Config(
        │   │   │   │   │   │   │           channels_in=96,
        │   │   │   │   │   │   │           channels_out=192,
        │   │   │   │   │   │   │           elements=[
        │   │   │   │   │   │   │               LayerNorm2d.Config(
        │   │   │   │   │   │   │                   channels_in=96,
        │   │   │   │   │   │   │                   channels_out=96,
        │   │   │   │   │   │   │                   eps=1e-06,
        │   │   │   │   │   │   │                   elementwise_affine=True,
        │   │   │   │   │   │   │                   device=None,
        │   │   │   │   │   │   │                   dtype=None
        │   │   │   │   │   │   │               ),
        │   │   │   │   │   │   │               Conv2d.Config(
        │   │   │   │   │   │   │                   channels_in=96,
        │   │   │   │   │   │   │                   channels_out=192,
        │   │   │   │   │   │   │                   kernel_size=2,
        │   │   │   │   │   │   │                   stride=2,
        │   │   │   │   │   │   │                   padding=0,
        │   │   │   │   │   │   │                   dilation=1,
        │   │   │   │   │   │   │                   groups=1,
        │   │   │   │   │   │   │                   bias=True,
        │   │   │   │   │   │   │                   depth_index=(),
        │   │   │   │   │   │   │                   device=None,
        │   │   │   │   │   │   │                   dtype=None,
        │   │   │   │   │   │   │                   init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │                   init_bias=None
        │   │   │   │   │   │   │               )
        │   │   │   │   │   │   │           ],
        │   │   │   │   │   │   │           repeat=1,
        │   │   │   │   │   │   │           depth_index=()
        │   │   │   │   │   │   │       ),
        │   │   │   │   │   │   │       block=[
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=192,
        │   │   │   │   │   │   │       │       channels_out=192,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=192,
        │   │   │   │   │   │   │       │           channels_out=192,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=192,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=192,
        │   │   │   │   │   │   │       │           channels_out=192,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=192,
        │   │   │   │   │   │   │       │           channels_out=192,
        │   │   │   │   │   │   │       │           channels_hidden=768,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=192,
        │   │   │   │   │   │   │       │               channels_out=768,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=768,
        │   │   │   │   │   │   │       │               channels_out=192,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=192,
        │   │   │   │   │   │   │       │       channels_out=192,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=192,
        │   │   │   │   │   │   │       │           channels_out=192,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=192,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=192,
        │   │   │   │   │   │   │       │           channels_out=192,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=192,
        │   │   │   │   │   │   │       │           channels_out=192,
        │   │   │   │   │   │   │       │           channels_hidden=768,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=192,
        │   │   │   │   │   │   │       │               channels_out=768,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=768,
        │   │   │   │   │   │   │       │               channels_out=192,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=192,
        │   │   │   │   │   │   │       │       channels_out=192,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=192,
        │   │   │   │   │   │   │       │           channels_out=192,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=192,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=192,
        │   │   │   │   │   │   │       │           channels_out=192,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=192,
        │   │   │   │   │   │   │       │           channels_out=192,
        │   │   │   │   │   │   │       │           channels_hidden=768,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=192,
        │   │   │   │   │   │   │       │               channels_out=768,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=768,
        │   │   │   │   │   │   │       │               channels_out=192,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   )
        │   │   │   │   │   │   │       ]
        │   │   │   │   │   │   │   ),
        │   │   │   │   │   │   │   ConvNeXtStage.Config(
        │   │   │   │   │   │   │       channels_in=192,
        │   │   │   │   │   │   │       channels_out=384,
        │   │   │   │   │   │   │       blocks_per_stage=2,
        │   │   │   │   │   │   │       downsample=Sequential.Config(
        │   │   │   │   │   │   │           channels_in=192,
        │   │   │   │   │   │   │           channels_out=384,
        │   │   │   │   │   │   │           elements=[
        │   │   │   │   │   │   │               LayerNorm2d.Config(
        │   │   │   │   │   │   │                   channels_in=192,
        │   │   │   │   │   │   │                   channels_out=192,
        │   │   │   │   │   │   │                   eps=1e-06,
        │   │   │   │   │   │   │                   elementwise_affine=True,
        │   │   │   │   │   │   │                   device=None,
        │   │   │   │   │   │   │                   dtype=None
        │   │   │   │   │   │   │               ),
        │   │   │   │   │   │   │               Conv2d.Config(
        │   │   │   │   │   │   │                   channels_in=192,
        │   │   │   │   │   │   │                   channels_out=384,
        │   │   │   │   │   │   │                   kernel_size=2,
        │   │   │   │   │   │   │                   stride=2,
        │   │   │   │   │   │   │                   padding=0,
        │   │   │   │   │   │   │                   dilation=1,
        │   │   │   │   │   │   │                   groups=1,
        │   │   │   │   │   │   │                   bias=True,
        │   │   │   │   │   │   │                   depth_index=(),
        │   │   │   │   │   │   │                   device=None,
        │   │   │   │   │   │   │                   dtype=None,
        │   │   │   │   │   │   │                   init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │                   init_bias=None
        │   │   │   │   │   │   │               )
        │   │   │   │   │   │   │           ],
        │   │   │   │   │   │   │           repeat=1,
        │   │   │   │   │   │   │           depth_index=()
        │   │   │   │   │   │   │       ),
        │   │   │   │   │   │   │       block=[
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=384,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │   │   │       │           channels_hidden=1_536,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=384,
        │   │   │   │   │   │   │       │               channels_out=1_536,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=1_536,
        │   │   │   │   │   │   │       │               channels_out=384,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   )
        │   │   │   │   │   │   │       ]
        │   │   │   │   │   │   │   ),
        │   │   │   │   │   │   │   ConvNeXtStage.Config(
        │   │   │   │   │   │   │       channels_in=384,
        │   │   │   │   │   │   │       channels_out=768,
        │   │   │   │   │   │   │       blocks_per_stage=2,
        │   │   │   │   │   │   │       downsample=Sequential.Config(
        │   │   │   │   │   │   │           channels_in=384,
        │   │   │   │   │   │   │           channels_out=768,
        │   │   │   │   │   │   │           elements=[
        │   │   │   │   │   │   │               LayerNorm2d.Config(
        │   │   │   │   │   │   │                   channels_in=384,
        │   │   │   │   │   │   │                   channels_out=384,
        │   │   │   │   │   │   │                   eps=1e-06,
        │   │   │   │   │   │   │                   elementwise_affine=True,
        │   │   │   │   │   │   │                   device=None,
        │   │   │   │   │   │   │                   dtype=None
        │   │   │   │   │   │   │               ),
        │   │   │   │   │   │   │               Conv2d.Config(
        │   │   │   │   │   │   │                   channels_in=384,
        │   │   │   │   │   │   │                   channels_out=768,
        │   │   │   │   │   │   │                   kernel_size=2,
        │   │   │   │   │   │   │                   stride=2,
        │   │   │   │   │   │   │                   padding=0,
        │   │   │   │   │   │   │                   dilation=1,
        │   │   │   │   │   │   │                   groups=1,
        │   │   │   │   │   │   │                   bias=True,
        │   │   │   │   │   │   │                   depth_index=(),
        │   │   │   │   │   │   │                   device=None,
        │   │   │   │   │   │   │                   dtype=None,
        │   │   │   │   │   │   │                   init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │                   init_bias=None
        │   │   │   │   │   │   │               )
        │   │   │   │   │   │   │           ],
        │   │   │   │   │   │   │           repeat=1,
        │   │   │   │   │   │   │           depth_index=()
        │   │   │   │   │   │   │       ),
        │   │   │   │   │   │   │       block=[
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=768,
        │   │   │   │   │   │   │       │       channels_out=768,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=768,
        │   │   │   │   │   │   │       │           channels_out=768,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=768,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=768,
        │   │   │   │   │   │   │       │           channels_out=768,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=768,
        │   │   │   │   │   │   │       │           channels_out=768,
        │   │   │   │   │   │   │       │           channels_hidden=3_072,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=768,
        │   │   │   │   │   │   │       │               channels_out=3_072,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=3_072,
        │   │   │   │   │   │   │       │               channels_out=768,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=768,
        │   │   │   │   │   │   │       │       channels_out=768,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=768,
        │   │   │   │   │   │   │       │           channels_out=768,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=768,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=768,
        │   │   │   │   │   │   │       │           channels_out=768,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=768,
        │   │   │   │   │   │   │       │           channels_out=768,
        │   │   │   │   │   │   │       │           channels_hidden=3_072,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=768,
        │   │   │   │   │   │   │       │               channels_out=3_072,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=3_072,
        │   │   │   │   │   │   │       │               channels_out=768,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   ),
        │   │   │   │   │   │   │       │   ConvNeXtBlock.Config(
        │   │   │   │   │   │   │       │       channels_in=768,
        │   │   │   │   │   │   │       │       channels_out=768,
        │   │   │   │   │   │   │       │       conv_dw=Conv2d.Config(
        │   │   │   │   │   │   │       │           channels_in=768,
        │   │   │   │   │   │   │       │           channels_out=768,
        │   │   │   │   │   │   │       │           kernel_size=7,
        │   │   │   │   │   │   │       │           stride=1,
        │   │   │   │   │   │   │       │           padding=3,
        │   │   │   │   │   │   │       │           dilation=1,
        │   │   │   │   │   │   │       │           groups=768,
        │   │   │   │   │   │   │       │           bias=True,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None,
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │           init_bias=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       norm=LayerNorm.Config(
        │   │   │   │   │   │   │       │           channels_in=768,
        │   │   │   │   │   │   │       │           channels_out=768,
        │   │   │   │   │   │   │       │           eps=1e-06,
        │   │   │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │   │   │       │           device=None,
        │   │   │   │   │   │   │       │           dtype=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       mlp=SwiGLU.Config(
        │   │   │   │   │   │   │       │           channels_in=768,
        │   │   │   │   │   │   │       │           channels_out=768,
        │   │   │   │   │   │   │       │           channels_hidden=3_072,
        │   │   │   │   │   │   │       │           expansion=4.0,
        │   │   │   │   │   │   │       │           round_to=1,
        │   │   │   │   │   │   │       │           gate=False,
        │   │   │   │   │   │   │       │           bias=False,
        │   │   │   │   │   │   │       │           split_gate_projection=False,
        │   │   │   │   │   │   │       │           norm=None,
        │   │   │   │   │   │   │       │           act=priml.model.swiglu.gelu,
        │   │   │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │   │   │       │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │   │       │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │   │       │           up_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=768,
        │   │   │   │   │   │   │       │               channels_out=3_072,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           middle=None,
        │   │   │   │   │   │   │       │           down_proj=Linear.Config(
        │   │   │   │   │   │   │       │               channels_in=3_072,
        │   │   │   │   │   │   │       │               channels_out=768,
        │   │   │   │   │   │   │       │               bias=True,
        │   │   │   │   │   │   │       │               depth_index=(),
        │   │   │   │   │   │   │       │               device=None,
        │   │   │   │   │   │   │       │               dtype=None,
        │   │   │   │   │   │   │       │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │   │       │               init_bias=None,
        │   │   │   │   │   │   │       │               linear_fn=<built-in function linear>,
        │   │   │   │   │   │   │       │               shard=None
        │   │   │   │   │   │   │       │           ),
        │   │   │   │   │   │   │       │           shard=None
        │   │   │   │   │   │   │       │       ),
        │   │   │   │   │   │   │       │       layer_scale=1e-06,
        │   │   │   │   │   │   │       │       drop_path=0.0
        │   │   │   │   │   │   │       │   )
        │   │   │   │   │   │   │       ]
        │   │   │   │   │   │   │   )
        │   │   │   │   │   │   ],
        │   │   │   │   │   │   head=ConvNeXtHead.Config(
        │   │   │   │   │   │       channels_in=768,
        │   │   │   │   │   │       norm=LayerNorm.Config(
        │   │   │   │   │   │           channels_in=768,
        │   │   │   │   │   │           channels_out=768,
        │   │   │   │   │   │           eps=1e-06,
        │   │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │   │           device=None,
        │   │   │   │   │   │           dtype=None
        │   │   │   │   │   │       ),
        │   │   │   │   │   │       proj_out=Identity.Config(
        │   │   │   │   │   │           channels_in=768,
        │   │   │   │   │   │           channels_out=768
        │   │   │   │   │   │       )
        │   │   │   │   │   │   ),
        │   │   │   │   │   │   drop_path=0.0,
        │   │   │   │   │   │   drop_path_schedule=priml.model.convnext.linear_drop_path
        │   │   │   │   │   ),
        │   │   │   │   │   checkpoint=PosixPath('/opt/scratch/datasets/rogii-ruby-fade6112/initialization/convnext_small.safetensors'),
        │   │   │   │   │   init_weight=functools.partial(torch.nn.init.trunc_normal_, std=0.02)
        │   │   │   │   ),
        │   │   │   │   up_blocks=[
        │   │   │   │   │   ConcatUp.Config(
        │   │   │   │   │       channels_in=768,
        │   │   │   │   │       channels_out=384,
        │   │   │   │   │       channels_skip=384,
        │   │   │   │   │       proj=Conv2d.Config(
        │   │   │   │   │           channels_in=1_152,
        │   │   │   │   │           channels_out=384,
        │   │   │   │   │           kernel_size=1,
        │   │   │   │   │           stride=1,
        │   │   │   │   │           padding='same',
        │   │   │   │   │           dilation=1,
        │   │   │   │   │           groups=1,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           depth_index=(),
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None,
        │   │   │   │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │           init_bias=torch.nn.init.zeros_
        │   │   │   │   │       ),
        │   │   │   │   │       block=ResidualBlock.Config(
        │   │   │   │   │       │   channels_in=384,
        │   │   │   │   │       │   channels_out=384,
        │   │   │   │   │       │   conv1=ConvNormAct.Config(
        │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │       │       conv=Conv2d.Config(
        │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │       │           kernel_size=3,
        │   │   │   │   │       │           stride=1,
        │   │   │   │   │       │           padding='same',
        │   │   │   │   │       │           dilation=1,
        │   │   │   │   │       │           groups=1,
        │   │   │   │   │       │           bias=False,
        │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │       │           device=None,
        │   │   │   │   │       │           dtype=None,
        │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │   │       │       ),
        │   │   │   │   │       │       norm=BatchNorm2d.Config(
        │   │   │   │   │       │           channels_in=384,
        │   │   │   │   │       │           channels_out=384,
        │   │   │   │   │       │           momentum=0.1,
        │   │   │   │   │       │           eps=1e-05,
        │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │       │           device=None,
        │   │   │   │   │       │           dtype=None
        │   │   │   │   │       │       ),
        │   │   │   │   │       │       activation=priml.model.swiglu.silu,
        │   │   │   │   │       │       drop=Dropout2d.Config(
        │   │   │   │   │       │           channels_in=-1,
        │   │   │   │   │       │           channels_out=-1,
        │   │   │   │   │       │           p=0.05,
        │   │   │   │   │       │           inplace=False
        │   │   │   │   │       │       )
        │   │   │   │   │       │   ),
        │   │   │   │   │       │   conv2=Conv2d.Config(
        │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │       │       kernel_size=3,
        │   │   │   │   │       │       stride=1,
        │   │   │   │   │       │       padding='same',
        │   │   │   │   │       │       dilation=1,
        │   │   │   │   │       │       groups=1,
        │   │   │   │   │       │       bias=False,
        │   │   │   │   │       │       depth_index=(),
        │   │   │   │   │       │       device=None,
        │   │   │   │   │       │       dtype=None,
        │   │   │   │   │       │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │       │       init_bias=torch.nn.init.zeros_
        │   │   │   │   │       │   ),
        │   │   │   │   │       │   norm=BatchNorm2d.Config(
        │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │       │       channels_out=384,
        │   │   │   │   │       │       momentum=0.1,
        │   │   │   │   │       │       eps=1e-05,
        │   │   │   │   │       │       elementwise_affine=True,
        │   │   │   │   │       │       device=None,
        │   │   │   │   │       │       dtype=None
        │   │   │   │   │       │   ),
        │   │   │   │   │       │   drop=Dropout2d.Config(
        │   │   │   │   │       │       channels_in=-1,
        │   │   │   │   │       │       channels_out=-1,
        │   │   │   │   │       │       p=0.05,
        │   │   │   │   │       │       inplace=False
        │   │   │   │   │       │   ),
        │   │   │   │   │       │   activation=priml.model.swiglu.silu,
        │   │   │   │   │       │   skip=Identity.Config(
        │   │   │   │   │       │       channels_in=384,
        │   │   │   │   │       │       channels_out=384
        │   │   │   │   │       │   )
        │   │   │   │   │       )
        │   │   │   │   │   ),
        │   │   │   │   │   ConcatUp.Config(
        │   │   │   │   │       channels_in=384,
        │   │   │   │   │       channels_out=192,
        │   │   │   │   │       channels_skip=192,
        │   │   │   │   │       proj=Conv2d.Config(
        │   │   │   │   │           channels_in=576,
        │   │   │   │   │           channels_out=192,
        │   │   │   │   │           kernel_size=1,
        │   │   │   │   │           stride=1,
        │   │   │   │   │           padding='same',
        │   │   │   │   │           dilation=1,
        │   │   │   │   │           groups=1,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           depth_index=(),
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None,
        │   │   │   │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │           init_bias=torch.nn.init.zeros_
        │   │   │   │   │       ),
        │   │   │   │   │       block=ResidualBlock.Config(
        │   │   │   │   │       │   channels_in=192,
        │   │   │   │   │       │   channels_out=192,
        │   │   │   │   │       │   conv1=ConvNormAct.Config(
        │   │   │   │   │       │       channels_in=192,
        │   │   │   │   │       │       channels_out=192,
        │   │   │   │   │       │       conv=Conv2d.Config(
        │   │   │   │   │       │           channels_in=192,
        │   │   │   │   │       │           channels_out=192,
        │   │   │   │   │       │           kernel_size=3,
        │   │   │   │   │       │           stride=1,
        │   │   │   │   │       │           padding='same',
        │   │   │   │   │       │           dilation=1,
        │   │   │   │   │       │           groups=1,
        │   │   │   │   │       │           bias=False,
        │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │       │           device=None,
        │   │   │   │   │       │           dtype=None,
        │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │   │       │       ),
        │   │   │   │   │       │       norm=BatchNorm2d.Config(
        │   │   │   │   │       │           channels_in=192,
        │   │   │   │   │       │           channels_out=192,
        │   │   │   │   │       │           momentum=0.1,
        │   │   │   │   │       │           eps=1e-05,
        │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │       │           device=None,
        │   │   │   │   │       │           dtype=None
        │   │   │   │   │       │       ),
        │   │   │   │   │       │       activation=priml.model.swiglu.silu,
        │   │   │   │   │       │       drop=Dropout2d.Config(
        │   │   │   │   │       │           channels_in=-1,
        │   │   │   │   │       │           channels_out=-1,
        │   │   │   │   │       │           p=0.05,
        │   │   │   │   │       │           inplace=False
        │   │   │   │   │       │       )
        │   │   │   │   │       │   ),
        │   │   │   │   │       │   conv2=Conv2d.Config(
        │   │   │   │   │       │       channels_in=192,
        │   │   │   │   │       │       channels_out=192,
        │   │   │   │   │       │       kernel_size=3,
        │   │   │   │   │       │       stride=1,
        │   │   │   │   │       │       padding='same',
        │   │   │   │   │       │       dilation=1,
        │   │   │   │   │       │       groups=1,
        │   │   │   │   │       │       bias=False,
        │   │   │   │   │       │       depth_index=(),
        │   │   │   │   │       │       device=None,
        │   │   │   │   │       │       dtype=None,
        │   │   │   │   │       │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │       │       init_bias=torch.nn.init.zeros_
        │   │   │   │   │       │   ),
        │   │   │   │   │       │   norm=BatchNorm2d.Config(
        │   │   │   │   │       │       channels_in=192,
        │   │   │   │   │       │       channels_out=192,
        │   │   │   │   │       │       momentum=0.1,
        │   │   │   │   │       │       eps=1e-05,
        │   │   │   │   │       │       elementwise_affine=True,
        │   │   │   │   │       │       device=None,
        │   │   │   │   │       │       dtype=None
        │   │   │   │   │       │   ),
        │   │   │   │   │       │   drop=Dropout2d.Config(
        │   │   │   │   │       │       channels_in=-1,
        │   │   │   │   │       │       channels_out=-1,
        │   │   │   │   │       │       p=0.05,
        │   │   │   │   │       │       inplace=False
        │   │   │   │   │       │   ),
        │   │   │   │   │       │   activation=priml.model.swiglu.silu,
        │   │   │   │   │       │   skip=Identity.Config(
        │   │   │   │   │       │       channels_in=192,
        │   │   │   │   │       │       channels_out=192
        │   │   │   │   │       │   )
        │   │   │   │   │       )
        │   │   │   │   │   ),
        │   │   │   │   │   ConcatUp.Config(
        │   │   │   │   │       channels_in=192,
        │   │   │   │   │       channels_out=96,
        │   │   │   │   │       channels_skip=96,
        │   │   │   │   │       proj=Conv2d.Config(
        │   │   │   │   │           channels_in=288,
        │   │   │   │   │           channels_out=96,
        │   │   │   │   │           kernel_size=1,
        │   │   │   │   │           stride=1,
        │   │   │   │   │           padding='same',
        │   │   │   │   │           dilation=1,
        │   │   │   │   │           groups=1,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           depth_index=(),
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None,
        │   │   │   │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │           init_bias=torch.nn.init.zeros_
        │   │   │   │   │       ),
        │   │   │   │   │       block=ResidualBlock.Config(
        │   │   │   │   │       │   channels_in=96,
        │   │   │   │   │       │   channels_out=96,
        │   │   │   │   │       │   conv1=ConvNormAct.Config(
        │   │   │   │   │       │       channels_in=96,
        │   │   │   │   │       │       channels_out=96,
        │   │   │   │   │       │       conv=Conv2d.Config(
        │   │   │   │   │       │           channels_in=96,
        │   │   │   │   │       │           channels_out=96,
        │   │   │   │   │       │           kernel_size=3,
        │   │   │   │   │       │           stride=1,
        │   │   │   │   │       │           padding='same',
        │   │   │   │   │       │           dilation=1,
        │   │   │   │   │       │           groups=1,
        │   │   │   │   │       │           bias=False,
        │   │   │   │   │       │           depth_index=(),
        │   │   │   │   │       │           device=None,
        │   │   │   │   │       │           dtype=None,
        │   │   │   │   │       │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │       │           init_bias=torch.nn.init.zeros_
        │   │   │   │   │       │       ),
        │   │   │   │   │       │       norm=BatchNorm2d.Config(
        │   │   │   │   │       │           channels_in=96,
        │   │   │   │   │       │           channels_out=96,
        │   │   │   │   │       │           momentum=0.1,
        │   │   │   │   │       │           eps=1e-05,
        │   │   │   │   │       │           elementwise_affine=True,
        │   │   │   │   │       │           device=None,
        │   │   │   │   │       │           dtype=None
        │   │   │   │   │       │       ),
        │   │   │   │   │       │       activation=priml.model.swiglu.silu,
        │   │   │   │   │       │       drop=Dropout2d.Config(
        │   │   │   │   │       │           channels_in=-1,
        │   │   │   │   │       │           channels_out=-1,
        │   │   │   │   │       │           p=0.05,
        │   │   │   │   │       │           inplace=False
        │   │   │   │   │       │       )
        │   │   │   │   │       │   ),
        │   │   │   │   │       │   conv2=Conv2d.Config(
        │   │   │   │   │       │       channels_in=96,
        │   │   │   │   │       │       channels_out=96,
        │   │   │   │   │       │       kernel_size=3,
        │   │   │   │   │       │       stride=1,
        │   │   │   │   │       │       padding='same',
        │   │   │   │   │       │       dilation=1,
        │   │   │   │   │       │       groups=1,
        │   │   │   │   │       │       bias=False,
        │   │   │   │   │       │       depth_index=(),
        │   │   │   │   │       │       device=None,
        │   │   │   │   │       │       dtype=None,
        │   │   │   │   │       │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │       │       init_bias=torch.nn.init.zeros_
        │   │   │   │   │       │   ),
        │   │   │   │   │       │   norm=BatchNorm2d.Config(
        │   │   │   │   │       │       channels_in=96,
        │   │   │   │   │       │       channels_out=96,
        │   │   │   │   │       │       momentum=0.1,
        │   │   │   │   │       │       eps=1e-05,
        │   │   │   │   │       │       elementwise_affine=True,
        │   │   │   │   │       │       device=None,
        │   │   │   │   │       │       dtype=None
        │   │   │   │   │       │   ),
        │   │   │   │   │       │   drop=Dropout2d.Config(
        │   │   │   │   │       │       channels_in=-1,
        │   │   │   │   │       │       channels_out=-1,
        │   │   │   │   │       │       p=0.05,
        │   │   │   │   │       │       inplace=False
        │   │   │   │   │       │   ),
        │   │   │   │   │       │   activation=priml.model.swiglu.silu,
        │   │   │   │   │       │   skip=Identity.Config(
        │   │   │   │   │       │       channels_in=96,
        │   │   │   │   │       │       channels_out=96
        │   │   │   │   │       │   )
        │   │   │   │   │       )
        │   │   │   │   │   )
        │   │   │   │   ],
        │   │   │   │   to_full_resolution=Conv2d.Config(
        │   │   │   │       channels_in=96,
        │   │   │   │       channels_out=32,
        │   │   │   │       kernel_size=1,
        │   │   │   │       stride=1,
        │   │   │   │       padding='same',
        │   │   │   │       dilation=1,
        │   │   │   │       groups=1,
        │   │   │   │       bias=False,
        │   │   │   │       depth_index=(),
        │   │   │   │       device=None,
        │   │   │   │       dtype=None,
        │   │   │   │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       init_bias=torch.nn.init.zeros_
        │   │   │   │   ),
        │   │   │   │   full_up_block=AddUp.Config(
        │   │   │   │   │   channels_in=32,
        │   │   │   │   │   block=ResidualBlock.Config(
        │   │   │   │   │   │   channels_in=32,
        │   │   │   │   │   │   channels_out=32,
        │   │   │   │   │   │   conv1=ConvNormAct.Config(
        │   │   │   │   │   │       channels_in=32,
        │   │   │   │   │   │       channels_out=32,
        │   │   │   │   │   │       conv=Conv2d.Config(
        │   │   │   │   │   │           channels_in=32,
        │   │   │   │   │   │           channels_out=32,
        │   │   │   │   │   │           kernel_size=3,
        │   │   │   │   │   │           stride=1,
        │   │   │   │   │   │           padding='same',
        │   │   │   │   │   │           dilation=1,
        │   │   │   │   │   │           groups=1,
        │   │   │   │   │   │           bias=False,
        │   │   │   │   │   │           depth_index=(),
        │   │   │   │   │   │           device=None,
        │   │   │   │   │   │           dtype=None,
        │   │   │   │   │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │           init_bias=torch.nn.init.zeros_
        │   │   │   │   │   │       ),
        │   │   │   │   │   │       norm=BatchNorm2d.Config(
        │   │   │   │   │   │           channels_in=32,
        │   │   │   │   │   │           channels_out=32,
        │   │   │   │   │   │           momentum=0.1,
        │   │   │   │   │   │           eps=1e-05,
        │   │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │   │           device=None,
        │   │   │   │   │   │           dtype=None
        │   │   │   │   │   │       ),
        │   │   │   │   │   │       activation=priml.model.swiglu.silu,
        │   │   │   │   │   │       drop=Dropout2d.Config(
        │   │   │   │   │   │           channels_in=-1,
        │   │   │   │   │   │           channels_out=-1,
        │   │   │   │   │   │           p=0.05,
        │   │   │   │   │   │           inplace=False
        │   │   │   │   │   │       )
        │   │   │   │   │   │   ),
        │   │   │   │   │   │   conv2=Conv2d.Config(
        │   │   │   │   │   │       channels_in=32,
        │   │   │   │   │   │       channels_out=32,
        │   │   │   │   │   │       kernel_size=3,
        │   │   │   │   │   │       stride=1,
        │   │   │   │   │   │       padding='same',
        │   │   │   │   │   │       dilation=1,
        │   │   │   │   │   │       groups=1,
        │   │   │   │   │   │       bias=False,
        │   │   │   │   │   │       depth_index=(),
        │   │   │   │   │   │       device=None,
        │   │   │   │   │   │       dtype=None,
        │   │   │   │   │   │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │       init_bias=torch.nn.init.zeros_
        │   │   │   │   │   │   ),
        │   │   │   │   │   │   norm=BatchNorm2d.Config(
        │   │   │   │   │   │       channels_in=32,
        │   │   │   │   │   │       channels_out=32,
        │   │   │   │   │   │       momentum=0.1,
        │   │   │   │   │   │       eps=1e-05,
        │   │   │   │   │   │       elementwise_affine=True,
        │   │   │   │   │   │       device=None,
        │   │   │   │   │   │       dtype=None
        │   │   │   │   │   │   ),
        │   │   │   │   │   │   drop=Dropout2d.Config(
        │   │   │   │   │   │       channels_in=-1,
        │   │   │   │   │   │       channels_out=-1,
        │   │   │   │   │   │       p=0.05,
        │   │   │   │   │   │       inplace=False
        │   │   │   │   │   │   ),
        │   │   │   │   │   │   activation=priml.model.swiglu.silu,
        │   │   │   │   │   │   skip=Identity.Config(
        │   │   │   │   │   │       channels_in=32,
        │   │   │   │   │   │       channels_out=32
        │   │   │   │   │   │   )
        │   │   │   │   │   )
        │   │   │   │   ),
        │   │   │   │   out=ResidualBlock.Config(
        │   │   │   │   │   channels_in=32,
        │   │   │   │   │   channels_out=32,
        │   │   │   │   │   conv1=ConvNormAct.Config(
        │   │   │   │   │       channels_in=32,
        │   │   │   │   │       channels_out=32,
        │   │   │   │   │       conv=Conv2d.Config(
        │   │   │   │   │           channels_in=32,
        │   │   │   │   │           channels_out=32,
        │   │   │   │   │           kernel_size=3,
        │   │   │   │   │           stride=1,
        │   │   │   │   │           padding='same',
        │   │   │   │   │           dilation=1,
        │   │   │   │   │           groups=1,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           depth_index=(),
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None,
        │   │   │   │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │           init_bias=torch.nn.init.zeros_
        │   │   │   │   │       ),
        │   │   │   │   │       norm=BatchNorm2d.Config(
        │   │   │   │   │           channels_in=32,
        │   │   │   │   │           channels_out=32,
        │   │   │   │   │           momentum=0.1,
        │   │   │   │   │           eps=1e-05,
        │   │   │   │   │           elementwise_affine=True,
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None
        │   │   │   │   │       ),
        │   │   │   │   │       activation=priml.model.swiglu.silu,
        │   │   │   │   │       drop=Dropout2d.Config(
        │   │   │   │   │           channels_in=-1,
        │   │   │   │   │           channels_out=-1,
        │   │   │   │   │           p=0.05,
        │   │   │   │   │           inplace=False
        │   │   │   │   │       )
        │   │   │   │   │   ),
        │   │   │   │   │   conv2=Conv2d.Config(
        │   │   │   │   │       channels_in=32,
        │   │   │   │   │       channels_out=32,
        │   │   │   │   │       kernel_size=3,
        │   │   │   │   │       stride=1,
        │   │   │   │   │       padding='same',
        │   │   │   │   │       dilation=1,
        │   │   │   │   │       groups=1,
        │   │   │   │   │       bias=False,
        │   │   │   │   │       depth_index=(),
        │   │   │   │   │       device=None,
        │   │   │   │   │       dtype=None,
        │   │   │   │   │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │       init_bias=torch.nn.init.zeros_
        │   │   │   │   │   ),
        │   │   │   │   │   norm=BatchNorm2d.Config(
        │   │   │   │   │       channels_in=32,
        │   │   │   │   │       channels_out=32,
        │   │   │   │   │       momentum=0.1,
        │   │   │   │   │       eps=1e-05,
        │   │   │   │   │       elementwise_affine=True,
        │   │   │   │   │       device=None,
        │   │   │   │   │       dtype=None
        │   │   │   │   │   ),
        │   │   │   │   │   drop=Dropout2d.Config(
        │   │   │   │   │       channels_in=-1,
        │   │   │   │   │       channels_out=-1,
        │   │   │   │   │       p=0.05,
        │   │   │   │   │       inplace=False
        │   │   │   │   │   ),
        │   │   │   │   │   activation=priml.model.swiglu.silu,
        │   │   │   │   │   skip=Identity.Config(channels_in=32, channels_out=32)
        │   │   │   │   )
        │   │   │   ),
        │   │   │   proj_out=Conv2d.Config(
        │   │   │       channels_in=32,
        │   │   │       channels_out=1,
        │   │   │       kernel_size=1,
        │   │   │       stride=1,
        │   │   │       padding='same',
        │   │   │       dilation=1,
        │   │   │       groups=1,
        │   │   │       bias=True,
        │   │   │       depth_index=(),
        │   │   │       device=None,
        │   │   │       dtype=None,
        │   │   │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │       init_bias=None
        │   │   │   ),
        │   │   │   unet_gr_rmse_head=Sequential.Config(
        │   │   │   │   channels_in=32,
        │   │   │   │   channels_out=1,
        │   │   │   │   elements=[
        │   │   │   │       Linear.Config(
        │   │   │   │           channels_in=32,
        │   │   │   │           channels_out=64,
        │   │   │   │           bias=True,
        │   │   │   │           depth_index=(),
        │   │   │   │           device=None,
        │   │   │   │           dtype=None,
        │   │   │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │           init_bias=None,
        │   │   │   │           linear_fn=<built-in function linear>,
        │   │   │   │           shard=None
        │   │   │   │       ),
        │   │   │   │       BatchNorm.Config(
        │   │   │   │           channels_in=64,
        │   │   │   │           channels_out=64,
        │   │   │   │           momentum=0.1,
        │   │   │   │           eps=1e-05,
        │   │   │   │           elementwise_affine=True,
        │   │   │   │           device=None,
        │   │   │   │           dtype=None
        │   │   │   │       ),
        │   │   │   │       Activation.Config(
        │   │   │   │           channels_in=64,
        │   │   │   │           channels_out=64,
        │   │   │   │           act=priml.model.swiglu.silu
        │   │   │   │       ),
        │   │   │   │       Dropout.Config(
        │   │   │   │           channels_in=64,
        │   │   │   │           channels_out=64,
        │   │   │   │           p=0.05,
        │   │   │   │           inplace=False
        │   │   │   │       ),
        │   │   │   │       Linear.Config(
        │   │   │   │           channels_in=64,
        │   │   │   │           channels_out=1,
        │   │   │   │           bias=True,
        │   │   │   │           depth_index=(),
        │   │   │   │           device=None,
        │   │   │   │           dtype=None,
        │   │   │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │           init_bias=None,
        │   │   │   │           linear_fn=<built-in function linear>,
        │   │   │   │           shard=None
        │   │   │   │       )
        │   │   │   │   ],
        │   │   │   │   repeat=1,
        │   │   │   │   depth_index=()
        │   │   │   )
        │   │   ),
        │   │   optimizer=CompositeOptimizer.Config(
        │   │       optimizers=[
        │   │           PartialConfig(
        │   │               <class 'torch.optim.adamw.AdamW'>,
        │   │               lr=0.0002,
        │   │               weight_decay=0.02
        │   │           )
        │   │       ],
        │   │       select=[],
        │   │       require_total=True,
        │   │       drop_empty=False
        │   │   ),
        │   │   loss=WellboreLoss.Config(
        │   │       objective=PosteriorLoss.Config(
        │   │           regression=PartialConfig(
        │   │               priml.baselines.wellbore_geology.loss.masked_huber,
        │   │               delta=1.0
        │   │           ),
        │   │           alignment=PartialConfig(
        │   │               priml.baselines.wellbore_geology.loss.masked_soft_cross_entropy
        │   │           ),
        │   │           gamma_ray=PartialConfig(
        │   │               priml.baselines.wellbore_geology.loss.masked_expectation
        │   │           ),
        │   │           regression_weight=0.016666666666666666,
        │   │           alignment_weight=0.06666666666666667,
        │   │           gamma_ray_weight=0.016666666666666666
        │   │       )
        │   │   ),
        │   │   accumulate_grad_batches=1,
        │   │   drop_partial_accumulation_on_epoch_end=True,
        │   │   train_budget_steps=inf,
        │   │   train_budget_sec=inf,
        │   │   train_budget_epochs=inf,
        │   │   lr_schedule=RecursiveCosine.Config(T_max=1, eta_min=5e-07),
        │   │   parallelism=NoParallel.Config(device=None),
        │   │   model_quantization=NoModelQuantization.Config(),
        │   │   activation_memoization=DefaultActivationStorage.Config(),
        │   │   compile=None,
        │   │   ema=EMA.Config(
        │   │       decay=0.99,
        │   │       update_after_step=0,
        │   │       update_every=1,
        │   │       track_buffers=True,
        │   │       shadow_kind='module',
        │   │       warmup_seed=False,
        │   │       decay_schedule=priml.train.ema.constant_decay,
        │   │       select=priml.optimizers.parameter_filter.everything,
        │   │       update_fn=priml.train.ema.average_floating_state,
        │   │       buffer_update_fn=priml.train.ema.average_persistent_buffers
        │   │   ),
        │   │   gradient_clip_norm=10.0,
        │   │   skip_step_on_nonfinite_grad=False,
        │   │   device_init='eager',
        │   │   dtype_autocast=torch.bfloat16,
        │   │   autocast_cache_enabled=True
        │   ),
        │   dataset=FoldData.Config(
        │   │   loader=WellLoader.Config(
        │   │   │   dataset=WellDataset.Config(
        │   │   │   │   observations=FoldObservations.Config(
        │   │   │   │       directory=PosixPath('/opt/scratch/datasets/rogii-ruby-fade6112/train'),
        │   │   │   │       support_directory=None,
        │   │   │   │       well_ids=(),
        │   │   │   │       support_well_ids=(),
        │   │   │   │       geography=GeographicPrior.Config(
        │   │   │   │           neighbor_wells=12,
        │   │   │   │           point_neighbors=112,
        │   │   │   │           train_bin_size=64,
        │   │   │   │           query_bin_size=32,
        │   │   │   │           prefix_bin_size=32,
        │   │   │   │           prefix_support_rows=512,
        │   │   │   │           surface_jump_limit=0.15,
        │   │   │   │           max_support_points=12_000,
        │   │   │   │           distance_epsilon=1e-06,
        │   │   │   │           distance_power=2.0,
        │   │   │   │           metric=AnisotropicMetric.Config(
        │   │   │   │               angle_degrees=50.0,
        │   │   │   │               ratio=1.3
        │   │   │   │           ),
        │   │   │   │           plane=RobustPlane.Config(
        │   │   │   │               ridge_fraction=0.01,
        │   │   │   │               min_eigenvalue_ratio=0.05,
        │   │   │   │               max_gradient=0.25,
        │   │   │   │               robust_scale=0.1,
        │   │   │   │               iterations=1
        │   │   │   │           )
        │   │   │   │       ),
        │   │   │   │       leave_one_out=True,
        │   │   │   │       calibration=ReferenceCalibration.Config(
        │   │   │   │           blend=0.7,
        │   │   │   │           raw_weight=0.3,
        │   │   │   │           window=100.0,
        │   │   │   │           resolution=4.0
        │   │   │   │       ),
        │   │   │   │       particle=None
        │   │   │   │   ),
        │   │   │   │   transform=AugmentedFeatures.Config(
        │   │   │   │   │   features=ImageFeatures.Config(
        │   │   │   │   │       prefix_rows=1_024,
        │   │   │   │   │       target_rows=10_000,
        │   │   │   │   │       downsample=32,
        │   │   │   │   │       reference_columns=400,
        │   │   │   │   │       reference_window=100.0,
        │   │   │   │   │       channels=[
        │   │   │   │   │           'tw_gr',
        │   │   │   │   │           'gr',
        │   │   │   │   │           'gr_abs_diff',
        │   │   │   │   │           'gr_isnan_rate',
        │   │   │   │   │           'gr_std',
        │   │   │   │   │           'gr_first_last_delta',
        │   │   │   │   │           'gr_slope',
        │   │   │   │   │           'gr_quadratic_a',
        │   │   │   │   │           'gr_quadratic_b',
        │   │   │   │   │           'gr_quadratic_c',
        │   │   │   │   │           'gr_quadratic_rmse',
        │   │   │   │   │           'tw_gr_is_nan',
        │   │   │   │   │           'tw_tvt_rel',
        │   │   │   │   │           'seen_tvt_rel',
        │   │   │   │   │           'tw_seen_tvt_abs_diff',
        │   │   │   │   │           'z_diff'
        │   │   │   │   │       ],
        │   │   │   │   │       particle=None,
        │   │   │   │   │       gamma_stats=(87.43, 21.09),
        │   │   │   │   │       reference_gamma_stats=(78.62, 28.81),
        │   │   │   │   │       seen_depth_stats=(-18.08, 48.21),
        │   │   │   │   │       target_stats=(1.591691, 15.853182),
        │   │   │   │   │       vertical_step_stats=(-0.016892, 0.1065),
        │   │   │   │   │       surface_step_stats=(0.005535, 0.037389),
        │   │   │   │   │       depth_axis_scale=57.73,
        │   │   │   │   │       alignment_sigma=1.25,
        │   │   │   │   │       gamma_penalty_clip=25.0,
        │   │   │   │   │       feature_clip=8.0,
        │   │   │   │   │       depth_difference_clip=100.0,
        │   │   │   │   │       quadratic_stats=(
        │   │   │   │   │           (0.000585828, 0.38249254),
        │   │   │   │   │           (0.001914565, 0.338615701),
        │   │   │   │   │           (-0.006968186, 1.025529919),
        │   │   │   │   │           (0.212819023, 0.131140903)
        │   │   │   │   │       )
        │   │   │   │   │   ),
        │   │   │   │   │   paths=PathSimulation.Config(
        │   │   │   │   │   │   trend=BootstrapTrend.Config(
        │   │   │   │   │   │       minimum_suffix_rows=1_000,
        │   │   │   │   │   │       maximum_step=0.15,
        │   │   │   │   │   │       blocks=[
        │   │   │   │   │   │           DerivativeBlock.Config(
        │   │   │   │   │   │               length=1_024,
        │   │   │   │   │   │               weight=0.45
        │   │   │   │   │   │           ),
        │   │   │   │   │   │           DerivativeBlock.Config(
        │   │   │   │   │   │               length=2_048,
        │   │   │   │   │   │               weight=0.4
        │   │   │   │   │   │           ),
        │   │   │   │   │   │           DerivativeBlock.Config(
        │   │   │   │   │   │               length=4_096,
        │   │   │   │   │   │               weight=0.15
        │   │   │   │   │   │           )
        │   │   │   │   │   │       ],
        │   │   │   │   │   │       join_rows=16,
        │   │   │   │   │   │       reverse_probability=0.0
        │   │   │   │   │   │   ),
        │   │   │   │   │   │   gamma_bias=GammaBias.Config(
        │   │   │   │   │   │       depth_std=0.95,
        │   │   │   │   │   │       components=[
        │   │   │   │   │   │           BoxcarComponent.Config(
        │   │   │   │   │   │               window=1_536,
        │   │   │   │   │   │               weight=0.68
        │   │   │   │   │   │           ),
        │   │   │   │   │   │           BoxcarComponent.Config(
        │   │   │   │   │   │               window=256,
        │   │   │   │   │   │               weight=0.2
        │   │   │   │   │   │           ),
        │   │   │   │   │   │           BoxcarComponent.Config(
        │   │   │   │   │   │               window=2,
        │   │   │   │   │   │               weight=0.12
        │   │   │   │   │   │           )
        │   │   │   │   │   │       ],
        │   │   │   │   │   │       ranges=[
        │   │   │   │   │   │           MismatchRange.Config(
        │   │   │   │   │   │               lower=5.6,
        │   │   │   │   │   │               upper=8.4,
        │   │   │   │   │   │               weight=0.24
        │   │   │   │   │   │           ),
        │   │   │   │   │   │           MismatchRange.Config(
        │   │   │   │   │   │               lower=8.4,
        │   │   │   │   │   │               upper=12.2,
        │   │   │   │   │   │               weight=0.51
        │   │   │   │   │   │           ),
        │   │   │   │   │   │           MismatchRange.Config(
        │   │   │   │   │   │               lower=12.2,
        │   │   │   │   │   │               upper=17.5,
        │   │   │   │   │   │               weight=0.25
        │   │   │   │   │   │           )
        │   │   │   │   │   │       ],
        │   │   │   │   │   │       offset_std=2.0,
        │   │   │   │   │   │       clip=36.0
        │   │   │   │   │   │   ),
        │   │   │   │   │   │   reference_jitter=ReferenceJitter.Config(
        │   │   │   │   │   │       probability=0.0,
        │   │   │   │   │   │       gain_std=0.04,
        │   │   │   │   │   │       gain_bounds=(0.9, 1.1),
        │   │   │   │   │   │       offset_std=2.0,
        │   │   │   │   │   │       drift_std=1.5,
        │   │   │   │   │   │       drift_windows=(15, 30, 60),
        │   │   │   │   │   │       smooth_probability=0.1,
        │   │   │   │   │   │       smooth_windows=(5, 11, 21),
        │   │   │   │   │   │       smooth_blend=(0.1, 0.35)
        │   │   │   │   │   │   ),
        │   │   │   │   │   │   reference_transform=ReferenceTransform.Config(
        │   │   │   │   │   │       probability=0.6,
        │   │   │   │   │   │       gain_std=0.03,
        │   │   │   │   │   │       offset_std=2.0,
        │   │   │   │   │   │       smooth_probability=0.1,
        │   │   │   │   │   │       sharpen_probability=0.1,
        │   │   │   │   │   │       smooth_window=5,
        │   │   │   │   │   │       sharpen_alpha=(0, 0.3)
        │   │   │   │   │   │   ),
        │   │   │   │   │   │   shift_probability=0.75,
        │   │   │   │   │   │   prefix_sample_rows=0,
        │   │   │   │   │   │   shift_nan_pattern=FixedNanPattern.Config(shift=False),
        │   │   │   │   │   │   shift_range=40.0,
        │   │   │   │   │   │   jump_threshold=0.15,
        │   │   │   │   │   │   fault_probability=0.05,
        │   │   │   │   │   │   maximum_faults=3,
        │   │   │   │   │   │   fault_std=1.75,
        │   │   │   │   │   │   reverse_probability=0.25,
        │   │   │   │   │   │   reverse_head_bound=10.0,
        │   │   │   │   │   │   reverse_visible_rows=256,
        │   │   │   │   │   │   gamma_noise_shift_probability=0.1,
        │   │   │   │   │   │   gamma_noise_shift_ratio=1.0,
        │   │   │   │   │   │   reference_drift_probability=0.0,
        │   │   │   │   │   │   reference_drift_std=1.0,
        │   │   │   │   │   │   stretch_probability=0.4,
        │   │   │   │   │   │   stretch_range=(0.7, 1.3),
        │   │   │   │   │   │   tail_cut_probability=0.2,
        │   │   │   │   │   │   maximum_tail_cut=0.5
        │   │   │   │   │   ),
        │   │   │   │   │   images=ImageAugmentation.Config(
        │   │   │   │   │       rotate_probability=0.1,
        │   │   │   │   │       ratio_jitter=0.25,
        │   │   │   │   │       maximum_shift=20.0,
        │   │   │   │   │       channel_probability=0.7,
        │   │   │   │   │       channel_mask_probability=0.1,
        │   │   │   │   │       particle_mask_probability=0.05,
        │   │   │   │   │       sequence_probability=0.1,
        │   │   │   │   │       sequence_mask_ratio=0.3
        │   │   │   │   │   )
        │   │   │   │   ),
        │   │   │   │   streams=WorkerStreams.Config(seed=-1),
        │   │   │   │   training=True
        │   │   │   ),
        │   │   │   seed=-1,
        │   │   │   batch_size=16,
        │   │   │   shuffle=True,
        │   │   │   drop_last=True,
        │   │   │   num_workers=4,
        │   │   │   prefetch_factor=1,
        │   │   │   persistent_workers=True,
        │   │   │   pin_memory=False
        │   │   ),
        │   │   validation=FoldObservations.Config(
        │   │       directory=PosixPath('/opt/scratch/datasets/rogii-ruby-fade6112/train'),
        │   │       support_directory=None,
        │   │       well_ids=(),
        │   │       support_well_ids=(),
        │   │       geography=GeographicPrior.Config(
        │   │           neighbor_wells=12,
        │   │           point_neighbors=112,
        │   │           train_bin_size=64,
        │   │           query_bin_size=32,
        │   │           prefix_bin_size=32,
        │   │           prefix_support_rows=512,
        │   │           surface_jump_limit=0.15,
        │   │           max_support_points=12_000,
        │   │           distance_epsilon=1e-06,
        │   │           distance_power=2.0,
        │   │           metric=AnisotropicMetric.Config(angle_degrees=50.0, ratio=1.3),
        │   │           plane=RobustPlane.Config(
        │   │               ridge_fraction=0.01,
        │   │               min_eigenvalue_ratio=0.05,
        │   │               max_gradient=0.25,
        │   │               robust_scale=0.1,
        │   │               iterations=1
        │   │           )
        │   │       ),
        │   │       leave_one_out=False,
        │   │       calibration=ReferenceCalibration.Config(
        │   │           blend=0.7,
        │   │           raw_weight=0.3,
        │   │           window=100.0,
        │   │           resolution=4.0
        │   │       ),
        │   │       particle=None
        │   │   ),
        │   │   validation_batch_size=32,
        │   │   validation_seed_offset=1,
        │   │   device=None
        │   ),
        │   metrics_train={},
        │   metrics_eval={   '': PooledRMSE.Config(
        │           expansion=ExpandPrediction.Config(
        │               prefix_rows=1_024,
        │               target_rows=10_000,
        │               downsample=32,
        │               target_stats=(1.591691, 15.853182)
        │           )
        │       )},
        │   checkpointer=None,
        │   profiler=None,
        │   phase_timer=PhaseTimer.Config(
        │       enabled=False,
        │       torch_profile=False,
        │       base_dir=PosixPath('/opt/scratch/opt/scratch/runs/wellbore_geology/exp000/0719_V1'),
        │       working_dir=PosixPath('/opt/scratch/opt/scratch/runs/wellbore_geology/exp000/0719_V1/profiling'),
        │       heartbeat_interval_sec=20.0,
        │       fault_dump_interval_sec=40.0,
        │       cuda_events=False
        │   ),
        │   tracker=None,
        │   max_steps=inf,
        │   max_epochs=300,
        │   max_time=inf,
        │   max_time_kind='wall',
        │   max_eval_time=inf,
        │   eval_stop_on_time_limit=False,
        │   num_steps_eval=0,
        │   num_steps_log=10,
        │   early_train_log_steps=0,
        │   eval_every_epoch=True,
        │   eval_extras_every_eval=False,
        │   eval_warmup_batches=0,
        │   eval_only=False,
        │   restore_rng_state=True,
        │   seed=None,
        │   model_torch_seed=torch.random.manual_seed,
        │   num_steps_garbage_collect=inf,
        │   mesh_dim_model_seed='pp',
        │   mesh_dim_data_seed='dp',
        │   runtime=CudnnSingleProcess.Config(
        │       device='cuda',
        │       deterministic=False,
        │       float32_matmul_precision='high',
        │       cudnn_benchmark=True,
        │       cudnn_deterministic=False,
        │       cudnn_allow_tf32=True
        │   ),
        │   eval_epochs=PartialConfig(
        │       priml.train.epoch_loop.first_periodic_and_last,
        │       every=5
        │   ),
        │   early_stopping=EarlyStopping.Config(
        │       metric='rmse',
        │       mode='min',
        │       patience=30,
        │       min_epochs=170
        │   ),
        │   best_metric='rmse',
        │   evaluate_truncated=False,
        │   fold=0,
        │   base_seed=-1,
        │   fold_seed_stride=1_000,
        │   loader_seed_offset=1
        ),
        cv=GeographicCV.Config(
            clusters=20,
            folds=5,
            repeats=3,
            cluster_seed=0,
            split_seed=33,
            extra_trials=2,
            tolerance=0.0001,
            max_iterations=300
        ),
        geography=GeographicPrior.Config(
            neighbor_wells=12,
            point_neighbors=112,
            train_bin_size=64,
            query_bin_size=32,
            prefix_bin_size=32,
            prefix_support_rows=512,
            surface_jump_limit=0.15,
            max_support_points=12_000,
            distance_epsilon=1e-06,
            distance_power=2.0,
            metric=AnisotropicMetric.Config(angle_degrees=50.0, ratio=1.3),
            plane=RobustPlane.Config(
                ridge_fraction=0.01,
                min_eigenvalue_ratio=0.05,
                max_gradient=0.25,
                robust_scale=0.1,
                iterations=1
            )
        ),
        particle=None,
        prediction=PredictionBuilder.Config(
            expansion=ExpandPrediction.Config(
                prefix_rows=1_024,
                target_rows=10_000,
                downsample=32,
                target_stats=(1.591691, 15.853182)
            ),
            smoothing=SurfaceSmoothing.Config(
                window=1_025,
                polyorder=1,
                blend=0.6828376753426914
            ),
            diagnostic_names=()
        ),
        inference=InferenceRecipe.Config(
            seed_offset=100_000,
            batch_size=-1,
            dtype_autocast=None,
            ema=NoEMA.Config(),
            parallelism=NoParallel.Config(device='cuda')
        ),
        dataset_directory=PosixPath('/opt/scratch/datasets/rogii-ruby-fade6112'),
        pf_directory=PosixPath('/opt/scratch/caches/wellbore_geology/exp000-pf/0719_V1'),
        prepare_particles=False,
        seed=7,
        first_n_folds=None
    )
