PublicRecipe.Config(
        study_name='casmi2026',
        experiment_name='exp000',
        base_dir='/opt/scratch',
        working_dir=PosixPath('/opt/scratch/runs/casmi2026/exp000'),
        dataset_dir=PosixPath('/opt/scratch/datasets/casmi2026'),
        queries=PosixPath('/opt/scratch/datasets/casmi2026/inputs/enveda-CASMI26-molecule-id-mass-spectra/test.parquet'),
        sample_csv=PosixPath('/opt/scratch/datasets/casmi2026/inputs/enveda-CASMI26-molecule-id-mass-spectra/sample_submission.csv'),
        input_dir=PosixPath('/opt/scratch/datasets/casmi2026/inputs'),
        dreams=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/hengck23/hengck23-dreams-enveda-casmi26'),
        protected_inputs=[PosixPath('/opt/scratch/datasets/casmi2026')],
        main=MainResource.Config(
        │   assets=MainAssets.Config(
        │   │   base_dir=PosixPath('/opt/scratch/datasets/casmi2026'),
        │   │   pool_dir=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/ahmedberatozer/casmi26-v2-pool'),
        │   │   manifest=LocalFile.Config(
        │   │       path=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/ahmedberatozer/casmi26-v4b-models/MANIFEST.json'),
        │   │       sha256='e64b960037769971e146a35932e29ed3f042de1716e4dbe281ebd84a7cfc3b6a'
        │   │   ),
        │   │   library_dir=PosixPath('/opt/scratch/datasets/casmi2026/training/work'),
        │   │   fingerprints=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/ahmedberatozer/casmi26-v2-pool/train_fp_sel.npy'),
        │   │   popularity_dir=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/dmitriigluzdov/casmi26-pubchem-popularity-prior'),
        │   │   chemistry=IdentityTable.Config(
        │   │       base_dir=PosixPath('/opt/scratch/datasets/casmi2026'),
        │   │       working_dir=PosixPath('/opt/scratch/datasets/casmi2026/native-chemistry/identities')
        │   │   ),
        │   │   reactions_dir=PosixPath('/opt/scratch/datasets/casmi2026/native-chemistry/reactions'),
        │   │   pool_structures_dir=PosixPath('/opt/scratch/datasets/casmi2026/native-chemistry/analog-pool'),
        │   │   train_structures_dir=PosixPath('/opt/scratch/datasets/casmi2026/native-chemistry/analog-train'),
        │   │   generated_structures_dir=PosixPath('/opt/scratch/datasets/casmi2026/native-chemistry/analog-generated'),
        │   │   pool_families_dir=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/ahmedberatozer/casmi26-v4b-models/code/fe_v4/res/frag2'),
        │   │   fragments_dir=PosixPath('/opt/scratch/datasets/casmi2026/native-chemistry/ahmed-fragments'),
        │   │   intrinsic_dir=PosixPath('/opt/scratch/datasets/casmi2026/native-chemistry/intrinsic'),
        │   │   fragnet_dir=PosixPath('/opt/scratch/datasets/casmi2026/native-chemistry/fragnet'),
        │   │   derivation_dir=PosixPath('/opt/scratch/datasets/casmi2026/native-chemistry/derivation'),
        │   │   priors=LocalFile.Config(
        │   │       path=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/ahmedberatozer/casmi26-v4b-models/code/fe_v4/res/derivation/prior_tables.pkl'),
        │   │       sha256='217905818dd187da47cb458a60aeae2a5f2b0139904b4e320f70e41a9bd74c90'
        │   │   ),
        │   │   stage_a=LocalFile.Config(
        │   │       path=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/ahmedberatozer/casmi26-v4b-models/code/fe_v4/res/fragnet/stageA.npz'),
        │   │       sha256='3a501ecbfae647d86d7fda87605eda2d573475201b2482ca02a130aa2a3d7a68'
        │   │   ),
        │   │   stage_b=LocalFile.Config(
        │   │       path=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/ahmedberatozer/casmi26-v4b-models/code/fe_v4/res/fragnet/stageB.txt'),
        │   │       sha256='6916c25c6f32fdea4a20f5c0b4b5cf45a3f6bbbfa625552195a7909e8ae8685b'
        │   │   ),
        │   │   rankers=LocalFile.Config(
        │   │       path=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/ahmedberatozer/casmi26-v4b-models/ranker_v4b.pkl'),
        │   │       sha256='6172e926fb2de244f0b9189ed54e1ed8a15e76d166592b724129a4d6a34931da'
        │   │   ),
        │   │   bank=MainBank.Config(
        │   │   │   checkpoints=[
        │   │   │   │   BankCheckpoint.Config(
        │   │   │   │       base_dir=PosixPath('/opt/scratch/datasets/casmi2026'),
        │   │   │   │       path=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/ahmedberatozer/casmi26-fpnet-full1/fpnet_full1.pt'),
        │   │   │   │       sha256='d13a3611e499d7783c4a6fce5b04363eb297dcf70f2d7d00bb66f9ff1f547123',
        │   │   │   │       model=FPNet.Config(
        │   │   │   │       │   channels_out=10_226,
        │   │   │   │       │   channels_hidden=512,
        │   │   │   │       │   num_layers=8,
        │   │   │   │       │   block=Block.Config(
        │   │   │   │       │   │   channels_in=512,
        │   │   │   │       │   │   channels_out=512,
        │   │   │   │       │   │   attn=PeakAttention.Config(
        │   │   │   │       │   │   │   channels_in=512,
        │   │   │   │       │   │   │   channels_out=512,
        │   │   │   │       │   │   │   attention=Attention.Config(
        │   │   │   │       │   │   │       channels_in=512,
        │   │   │   │       │   │   │       channels_out=512,
        │   │   │   │       │   │   │       num_heads=8,
        │   │   │   │       │   │   │       channels_head=64,
        │   │   │   │       │   │   │       num_heads_kv=8,
        │   │   │   │       │   │   │       bias=False,
        │   │   │   │       │   │   │       dropout=0.0,
        │   │   │   │       │   │   │       causal=False,
        │   │   │   │       │   │   │       rope=None,
        │   │   │   │       │   │   │       rotate=None,
        │   │   │   │       │   │   │       norm_qk=None,
        │   │   │   │       │   │   │       share_qk_norm=True,
        │   │   │   │       │   │   │       norm_out=None,
        │   │   │   │       │   │   │       split_qkv_projection=False,
        │   │   │   │       │   │   │       proj_qkv=Linear.Config(
        │   │   │   │       │   │   │           channels_in=512,
        │   │   │   │       │   │   │           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
        │   │   │   │       │   │   │       ),
        │   │   │   │       │   │   │       proj_out=Linear.Config(
        │   │   │   │       │   │   │           channels_in=512,
        │   │   │   │       │   │   │           channels_out=512,
        │   │   │   │       │   │   │           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
        │   │   │   │       │   │   │       ),
        │   │   │   │       │   │   │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │       depth_index=(),
        │   │   │   │       │   │   │       attn_kernel=SdpaFused.Config()
        │   │   │   │       │   │   │   ),
        │   │   │   │       │   │   │   drop=Dropout.Config(
        │   │   │   │       │   │   │       channels_in=512,
        │   │   │   │       │   │   │       channels_out=512,
        │   │   │   │       │   │   │       p=0.1,
        │   │   │   │       │   │   │       inplace=False
        │   │   │   │       │   │   │   )
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   ffn=Sequential.Config(
        │   │   │   │       │   │   │   channels_in=512,
        │   │   │   │       │   │   │   channels_out=512,
        │   │   │   │       │   │   │   elements=[
        │   │   │   │       │   │   │   │   SwiGLU.Config(
        │   │   │   │       │   │   │   │       channels_in=512,
        │   │   │   │       │   │   │   │       channels_out=512,
        │   │   │   │       │   │   │   │       channels_hidden=2_048,
        │   │   │   │       │   │   │   │       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=512,
        │   │   │   │       │   │   │   │           channels_out=2_048,
        │   │   │   │       │   │   │   │           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=Dropout.Config(
        │   │   │   │       │   │   │   │           channels_in=2_048,
        │   │   │   │       │   │   │   │           channels_out=2_048,
        │   │   │   │       │   │   │   │           p=0.1,
        │   │   │   │       │   │   │   │           inplace=False
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       down_proj=Linear.Config(
        │   │   │   │       │   │   │   │           channels_in=2_048,
        │   │   │   │       │   │   │   │           channels_out=512,
        │   │   │   │       │   │   │   │           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
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   Dropout.Config(
        │   │   │   │       │   │   │   │       channels_in=512,
        │   │   │   │       │   │   │   │       channels_out=512,
        │   │   │   │       │   │   │   │       p=0.1,
        │   │   │   │       │   │   │   │       inplace=False
        │   │   │   │       │   │   │   │   )
        │   │   │   │       │   │   │   ],
        │   │   │   │       │   │   │   repeat=1,
        │   │   │   │       │   │   │   depth_index=()
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   norm1=LayerNorm.Config(
        │   │   │   │       │   │       channels_in=512,
        │   │   │   │       │   │       channels_out=512,
        │   │   │   │       │   │       eps=1e-05,
        │   │   │   │       │   │       elementwise_affine=True,
        │   │   │   │       │   │       device=None,
        │   │   │   │       │   │       dtype=None
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   norm2=LayerNorm.Config(
        │   │   │   │       │   │       channels_in=512,
        │   │   │   │       │   │       channels_out=512,
        │   │   │   │       │   │       eps=1e-05,
        │   │   │   │       │   │       elementwise_affine=True,
        │   │   │   │       │   │       device=None,
        │   │   │   │       │   │       dtype=None
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   prenorm=True,
        │   │   │   │       │   │   checkpoint=False,
        │   │   │   │       │   │   depth_index=()
        │   │   │   │       │   ),
        │   │   │   │       │   mz_emb=SinEmb.Config(
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       lo=-2.0,
        │   │   │   │       │       hi=3.3,
        │   │   │   │       │       power=1.0
        │   │   │   │       │   ),
        │   │   │   │       │   nl_emb=SinEmb.Config(
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       lo=-2.0,
        │   │   │   │       │       hi=3.3,
        │   │   │   │       │       power=1.0
        │   │   │   │       │   ),
        │   │   │   │       │   prec_emb=SinEmb.Config(
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       lo=-2.0,
        │   │   │   │       │       hi=3.3,
        │   │   │   │       │       power=1.0
        │   │   │   │       │   ),
        │   │   │   │       │   ce_emb=SinEmb.Config(
        │   │   │   │       │       channels_out=128,
        │   │   │   │       │       lo=0.0,
        │   │   │   │       │       hi=2.2,
        │   │   │   │       │       power=1.0
        │   │   │   │       │   ),
        │   │   │   │       │   int_emb=Linear.Config(
        │   │   │   │       │       channels_in=2,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       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
        │   │   │   │       │   ),
        │   │   │   │       │   pk=Linear.Config(
        │   │   │   │       │       channels_in=1_536,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       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
        │   │   │   │       │   ),
        │   │   │   │       │   adduct=Embedding.Config(
        │   │   │   │       │       channels_in=60,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       padding_idx=None,
        │   │   │   │       │       device=None,
        │   │   │   │       │       dtype=None,
        │   │   │   │       │       lookup_dtype=None,
        │   │   │   │       │       shard=None,
        │   │   │   │       │       depth_index=(),
        │   │   │   │       │       init_weight=torch.nn.init.normal_
        │   │   │   │       │   ),
        │   │   │   │       │   gl=Linear.Config(
        │   │   │   │       │       channels_in=644,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       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
        │   │   │   │       │   ),
        │   │   │   │       │   norm=LayerNorm.Config(
        │   │   │   │       │       channels_in=512,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       eps=1e-05,
        │   │   │   │       │       elementwise_affine=True,
        │   │   │   │       │       device=None,
        │   │   │   │       │       dtype=None
        │   │   │   │       │   ),
        │   │   │   │       │   head=SwiGLU.Config(
        │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │       channels_out=10_226,
        │   │   │   │       │       channels_hidden=2_048,
        │   │   │   │       │       expansion=4.0,
        │   │   │   │       │       round_to=256,
        │   │   │   │       │       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=1_024,
        │   │   │   │       │           channels_out=2_048,
        │   │   │   │       │           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=Dropout.Config(
        │   │   │   │       │           channels_in=2_048,
        │   │   │   │       │           channels_out=2_048,
        │   │   │   │       │           p=0.1,
        │   │   │   │       │           inplace=False
        │   │   │   │       │       ),
        │   │   │   │       │       down_proj=Linear.Config(
        │   │   │   │       │           channels_in=2_048,
        │   │   │   │       │           channels_out=10_226,
        │   │   │   │       │           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
        │   │   │   │       │   ),
        │   │   │   │       │   head_elements=SwiGLU.Config(
        │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │       channels_out=10,
        │   │   │   │       │       channels_hidden=256,
        │   │   │   │       │       expansion=4.0,
        │   │   │   │       │       round_to=256,
        │   │   │   │       │       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=1_024,
        │   │   │   │       │           channels_out=256,
        │   │   │   │       │           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=256,
        │   │   │   │       │           channels_out=10,
        │   │   │   │       │           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
        │   │   │   │       │   )
        │   │   │   │       ),
        │   │   │   │       dtype_autocast=torch.bfloat16
        │   │   │   │   )
        │   │   │   ],
        │   │   │   device='cuda',
        │   │   │   batch_size=64,
        │   │   │   max_peaks=160,
        │   │   │   floor=0.001,
        │   │   │   merge_tolerance=0.005,
        │   │   │   merged_adduct=priml.baselines.casmi2026.dreams.modal_adduct,
        │   │   │   single_count=priml.baselines.casmi2026.dreams.keep_count
        │   │   ),
        │   │   views_bank=MainBank.Config(
        │   │   │   checkpoints=[
        │   │   │   │   BankCheckpoint.Config(
        │   │   │   │       base_dir=PosixPath('/opt/scratch/datasets/casmi2026'),
        │   │   │   │       path=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/ahmedberatozer/casmi26-v4b-models/fe_A.pt'),
        │   │   │   │       sha256='edccf479a2a2e53461a06bc5e1aee4341b22abb32302c4f6d2c45281000d742a',
        │   │   │   │       model=FPNet.Config(
        │   │   │   │       │   channels_out=10_226,
        │   │   │   │       │   channels_hidden=512,
        │   │   │   │       │   num_layers=8,
        │   │   │   │       │   block=Block.Config(
        │   │   │   │       │   │   channels_in=512,
        │   │   │   │       │   │   channels_out=512,
        │   │   │   │       │   │   attn=PeakAttention.Config(
        │   │   │   │       │   │   │   channels_in=512,
        │   │   │   │       │   │   │   channels_out=512,
        │   │   │   │       │   │   │   attention=Attention.Config(
        │   │   │   │       │   │   │       channels_in=512,
        │   │   │   │       │   │   │       channels_out=512,
        │   │   │   │       │   │   │       num_heads=8,
        │   │   │   │       │   │   │       channels_head=64,
        │   │   │   │       │   │   │       num_heads_kv=8,
        │   │   │   │       │   │   │       bias=False,
        │   │   │   │       │   │   │       dropout=0.0,
        │   │   │   │       │   │   │       causal=False,
        │   │   │   │       │   │   │       rope=None,
        │   │   │   │       │   │   │       rotate=None,
        │   │   │   │       │   │   │       norm_qk=None,
        │   │   │   │       │   │   │       share_qk_norm=True,
        │   │   │   │       │   │   │       norm_out=None,
        │   │   │   │       │   │   │       split_qkv_projection=False,
        │   │   │   │       │   │   │       proj_qkv=Linear.Config(
        │   │   │   │       │   │   │           channels_in=512,
        │   │   │   │       │   │   │           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
        │   │   │   │       │   │   │       ),
        │   │   │   │       │   │   │       proj_out=Linear.Config(
        │   │   │   │       │   │   │           channels_in=512,
        │   │   │   │       │   │   │           channels_out=512,
        │   │   │   │       │   │   │           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
        │   │   │   │       │   │   │       ),
        │   │   │   │       │   │   │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │       depth_index=(),
        │   │   │   │       │   │   │       attn_kernel=SdpaFused.Config()
        │   │   │   │       │   │   │   ),
        │   │   │   │       │   │   │   drop=Dropout.Config(
        │   │   │   │       │   │   │       channels_in=512,
        │   │   │   │       │   │   │       channels_out=512,
        │   │   │   │       │   │   │       p=0.1,
        │   │   │   │       │   │   │       inplace=False
        │   │   │   │       │   │   │   )
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   ffn=Sequential.Config(
        │   │   │   │       │   │   │   channels_in=512,
        │   │   │   │       │   │   │   channels_out=512,
        │   │   │   │       │   │   │   elements=[
        │   │   │   │       │   │   │   │   SwiGLU.Config(
        │   │   │   │       │   │   │   │       channels_in=512,
        │   │   │   │       │   │   │   │       channels_out=512,
        │   │   │   │       │   │   │   │       channels_hidden=2_048,
        │   │   │   │       │   │   │   │       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=512,
        │   │   │   │       │   │   │   │           channels_out=2_048,
        │   │   │   │       │   │   │   │           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=Dropout.Config(
        │   │   │   │       │   │   │   │           channels_in=2_048,
        │   │   │   │       │   │   │   │           channels_out=2_048,
        │   │   │   │       │   │   │   │           p=0.1,
        │   │   │   │       │   │   │   │           inplace=False
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       down_proj=Linear.Config(
        │   │   │   │       │   │   │   │           channels_in=2_048,
        │   │   │   │       │   │   │   │           channels_out=512,
        │   │   │   │       │   │   │   │           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
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   Dropout.Config(
        │   │   │   │       │   │   │   │       channels_in=512,
        │   │   │   │       │   │   │   │       channels_out=512,
        │   │   │   │       │   │   │   │       p=0.1,
        │   │   │   │       │   │   │   │       inplace=False
        │   │   │   │       │   │   │   │   )
        │   │   │   │       │   │   │   ],
        │   │   │   │       │   │   │   repeat=1,
        │   │   │   │       │   │   │   depth_index=()
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   norm1=LayerNorm.Config(
        │   │   │   │       │   │       channels_in=512,
        │   │   │   │       │   │       channels_out=512,
        │   │   │   │       │   │       eps=1e-05,
        │   │   │   │       │   │       elementwise_affine=True,
        │   │   │   │       │   │       device=None,
        │   │   │   │       │   │       dtype=None
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   norm2=LayerNorm.Config(
        │   │   │   │       │   │       channels_in=512,
        │   │   │   │       │   │       channels_out=512,
        │   │   │   │       │   │       eps=1e-05,
        │   │   │   │       │   │       elementwise_affine=True,
        │   │   │   │       │   │       device=None,
        │   │   │   │       │   │       dtype=None
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   prenorm=True,
        │   │   │   │       │   │   checkpoint=False,
        │   │   │   │       │   │   depth_index=()
        │   │   │   │       │   ),
        │   │   │   │       │   mz_emb=SinEmb.Config(
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       lo=-2.0,
        │   │   │   │       │       hi=3.3,
        │   │   │   │       │       power=1.0
        │   │   │   │       │   ),
        │   │   │   │       │   nl_emb=SinEmb.Config(
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       lo=-2.0,
        │   │   │   │       │       hi=3.3,
        │   │   │   │       │       power=1.0
        │   │   │   │       │   ),
        │   │   │   │       │   prec_emb=SinEmb.Config(
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       lo=-2.0,
        │   │   │   │       │       hi=3.3,
        │   │   │   │       │       power=1.0
        │   │   │   │       │   ),
        │   │   │   │       │   ce_emb=SinEmb.Config(
        │   │   │   │       │       channels_out=128,
        │   │   │   │       │       lo=0.0,
        │   │   │   │       │       hi=2.2,
        │   │   │   │       │       power=1.0
        │   │   │   │       │   ),
        │   │   │   │       │   int_emb=Linear.Config(
        │   │   │   │       │       channels_in=2,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       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
        │   │   │   │       │   ),
        │   │   │   │       │   pk=Linear.Config(
        │   │   │   │       │       channels_in=1_536,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       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
        │   │   │   │       │   ),
        │   │   │   │       │   adduct=Embedding.Config(
        │   │   │   │       │       channels_in=60,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       padding_idx=None,
        │   │   │   │       │       device=None,
        │   │   │   │       │       dtype=None,
        │   │   │   │       │       lookup_dtype=None,
        │   │   │   │       │       shard=None,
        │   │   │   │       │       depth_index=(),
        │   │   │   │       │       init_weight=torch.nn.init.normal_
        │   │   │   │       │   ),
        │   │   │   │       │   gl=Linear.Config(
        │   │   │   │       │       channels_in=644,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       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
        │   │   │   │       │   ),
        │   │   │   │       │   norm=LayerNorm.Config(
        │   │   │   │       │       channels_in=512,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       eps=1e-05,
        │   │   │   │       │       elementwise_affine=True,
        │   │   │   │       │       device=None,
        │   │   │   │       │       dtype=None
        │   │   │   │       │   ),
        │   │   │   │       │   head=SwiGLU.Config(
        │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │       channels_out=10_226,
        │   │   │   │       │       channels_hidden=2_048,
        │   │   │   │       │       expansion=4.0,
        │   │   │   │       │       round_to=256,
        │   │   │   │       │       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=1_024,
        │   │   │   │       │           channels_out=2_048,
        │   │   │   │       │           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=Dropout.Config(
        │   │   │   │       │           channels_in=2_048,
        │   │   │   │       │           channels_out=2_048,
        │   │   │   │       │           p=0.1,
        │   │   │   │       │           inplace=False
        │   │   │   │       │       ),
        │   │   │   │       │       down_proj=Linear.Config(
        │   │   │   │       │           channels_in=2_048,
        │   │   │   │       │           channels_out=10_226,
        │   │   │   │       │           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
        │   │   │   │       │   ),
        │   │   │   │       │   head_elements=SwiGLU.Config(
        │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │       channels_out=10,
        │   │   │   │       │       channels_hidden=256,
        │   │   │   │       │       expansion=4.0,
        │   │   │   │       │       round_to=256,
        │   │   │   │       │       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=1_024,
        │   │   │   │       │           channels_out=256,
        │   │   │   │       │           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=256,
        │   │   │   │       │           channels_out=10,
        │   │   │   │       │           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
        │   │   │   │       │   )
        │   │   │   │       ),
        │   │   │   │       dtype_autocast=torch.bfloat16
        │   │   │   │   ),
        │   │   │   │   BankCheckpoint.Config(
        │   │   │   │       base_dir=PosixPath('/opt/scratch/datasets/casmi2026'),
        │   │   │   │       path=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/ahmedberatozer/casmi26-v4b-models/fe_B.pt'),
        │   │   │   │       sha256='6f60a22b3c3b21aa48d47328bb9523bdf07cc004bfb72966e0f9713507c24f2e',
        │   │   │   │       model=FPNet.Config(
        │   │   │   │       │   channels_out=10_226,
        │   │   │   │       │   channels_hidden=512,
        │   │   │   │       │   num_layers=8,
        │   │   │   │       │   block=Block.Config(
        │   │   │   │       │   │   channels_in=512,
        │   │   │   │       │   │   channels_out=512,
        │   │   │   │       │   │   attn=PeakAttention.Config(
        │   │   │   │       │   │   │   channels_in=512,
        │   │   │   │       │   │   │   channels_out=512,
        │   │   │   │       │   │   │   attention=Attention.Config(
        │   │   │   │       │   │   │       channels_in=512,
        │   │   │   │       │   │   │       channels_out=512,
        │   │   │   │       │   │   │       num_heads=8,
        │   │   │   │       │   │   │       channels_head=64,
        │   │   │   │       │   │   │       num_heads_kv=8,
        │   │   │   │       │   │   │       bias=False,
        │   │   │   │       │   │   │       dropout=0.0,
        │   │   │   │       │   │   │       causal=False,
        │   │   │   │       │   │   │       rope=None,
        │   │   │   │       │   │   │       rotate=None,
        │   │   │   │       │   │   │       norm_qk=None,
        │   │   │   │       │   │   │       share_qk_norm=True,
        │   │   │   │       │   │   │       norm_out=None,
        │   │   │   │       │   │   │       split_qkv_projection=False,
        │   │   │   │       │   │   │       proj_qkv=Linear.Config(
        │   │   │   │       │   │   │           channels_in=512,
        │   │   │   │       │   │   │           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
        │   │   │   │       │   │   │       ),
        │   │   │   │       │   │   │       proj_out=Linear.Config(
        │   │   │   │       │   │   │           channels_in=512,
        │   │   │   │       │   │   │           channels_out=512,
        │   │   │   │       │   │   │           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
        │   │   │   │       │   │   │       ),
        │   │   │   │       │   │   │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │       depth_index=(),
        │   │   │   │       │   │   │       attn_kernel=SdpaFused.Config()
        │   │   │   │       │   │   │   ),
        │   │   │   │       │   │   │   drop=Dropout.Config(
        │   │   │   │       │   │   │       channels_in=512,
        │   │   │   │       │   │   │       channels_out=512,
        │   │   │   │       │   │   │       p=0.1,
        │   │   │   │       │   │   │       inplace=False
        │   │   │   │       │   │   │   )
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   ffn=Sequential.Config(
        │   │   │   │       │   │   │   channels_in=512,
        │   │   │   │       │   │   │   channels_out=512,
        │   │   │   │       │   │   │   elements=[
        │   │   │   │       │   │   │   │   SwiGLU.Config(
        │   │   │   │       │   │   │   │       channels_in=512,
        │   │   │   │       │   │   │   │       channels_out=512,
        │   │   │   │       │   │   │   │       channels_hidden=2_048,
        │   │   │   │       │   │   │   │       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=512,
        │   │   │   │       │   │   │   │           channels_out=2_048,
        │   │   │   │       │   │   │   │           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=Dropout.Config(
        │   │   │   │       │   │   │   │           channels_in=2_048,
        │   │   │   │       │   │   │   │           channels_out=2_048,
        │   │   │   │       │   │   │   │           p=0.1,
        │   │   │   │       │   │   │   │           inplace=False
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       down_proj=Linear.Config(
        │   │   │   │       │   │   │   │           channels_in=2_048,
        │   │   │   │       │   │   │   │           channels_out=512,
        │   │   │   │       │   │   │   │           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
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   Dropout.Config(
        │   │   │   │       │   │   │   │       channels_in=512,
        │   │   │   │       │   │   │   │       channels_out=512,
        │   │   │   │       │   │   │   │       p=0.1,
        │   │   │   │       │   │   │   │       inplace=False
        │   │   │   │       │   │   │   │   )
        │   │   │   │       │   │   │   ],
        │   │   │   │       │   │   │   repeat=1,
        │   │   │   │       │   │   │   depth_index=()
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   norm1=LayerNorm.Config(
        │   │   │   │       │   │       channels_in=512,
        │   │   │   │       │   │       channels_out=512,
        │   │   │   │       │   │       eps=1e-05,
        │   │   │   │       │   │       elementwise_affine=True,
        │   │   │   │       │   │       device=None,
        │   │   │   │       │   │       dtype=None
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   norm2=LayerNorm.Config(
        │   │   │   │       │   │       channels_in=512,
        │   │   │   │       │   │       channels_out=512,
        │   │   │   │       │   │       eps=1e-05,
        │   │   │   │       │   │       elementwise_affine=True,
        │   │   │   │       │   │       device=None,
        │   │   │   │       │   │       dtype=None
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   prenorm=True,
        │   │   │   │       │   │   checkpoint=False,
        │   │   │   │       │   │   depth_index=()
        │   │   │   │       │   ),
        │   │   │   │       │   mz_emb=SinEmb.Config(
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       lo=-2.0,
        │   │   │   │       │       hi=3.3,
        │   │   │   │       │       power=1.0
        │   │   │   │       │   ),
        │   │   │   │       │   nl_emb=SinEmb.Config(
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       lo=-2.0,
        │   │   │   │       │       hi=3.3,
        │   │   │   │       │       power=1.0
        │   │   │   │       │   ),
        │   │   │   │       │   prec_emb=SinEmb.Config(
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       lo=-2.0,
        │   │   │   │       │       hi=3.3,
        │   │   │   │       │       power=1.0
        │   │   │   │       │   ),
        │   │   │   │       │   ce_emb=SinEmb.Config(
        │   │   │   │       │       channels_out=128,
        │   │   │   │       │       lo=0.0,
        │   │   │   │       │       hi=2.2,
        │   │   │   │       │       power=1.0
        │   │   │   │       │   ),
        │   │   │   │       │   int_emb=Linear.Config(
        │   │   │   │       │       channels_in=2,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       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
        │   │   │   │       │   ),
        │   │   │   │       │   pk=Linear.Config(
        │   │   │   │       │       channels_in=1_536,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       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
        │   │   │   │       │   ),
        │   │   │   │       │   adduct=Embedding.Config(
        │   │   │   │       │       channels_in=60,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       padding_idx=None,
        │   │   │   │       │       device=None,
        │   │   │   │       │       dtype=None,
        │   │   │   │       │       lookup_dtype=None,
        │   │   │   │       │       shard=None,
        │   │   │   │       │       depth_index=(),
        │   │   │   │       │       init_weight=torch.nn.init.normal_
        │   │   │   │       │   ),
        │   │   │   │       │   gl=Linear.Config(
        │   │   │   │       │       channels_in=644,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       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
        │   │   │   │       │   ),
        │   │   │   │       │   norm=LayerNorm.Config(
        │   │   │   │       │       channels_in=512,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       eps=1e-05,
        │   │   │   │       │       elementwise_affine=True,
        │   │   │   │       │       device=None,
        │   │   │   │       │       dtype=None
        │   │   │   │       │   ),
        │   │   │   │       │   head=SwiGLU.Config(
        │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │       channels_out=10_226,
        │   │   │   │       │       channels_hidden=2_048,
        │   │   │   │       │       expansion=4.0,
        │   │   │   │       │       round_to=256,
        │   │   │   │       │       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=1_024,
        │   │   │   │       │           channels_out=2_048,
        │   │   │   │       │           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=Dropout.Config(
        │   │   │   │       │           channels_in=2_048,
        │   │   │   │       │           channels_out=2_048,
        │   │   │   │       │           p=0.1,
        │   │   │   │       │           inplace=False
        │   │   │   │       │       ),
        │   │   │   │       │       down_proj=Linear.Config(
        │   │   │   │       │           channels_in=2_048,
        │   │   │   │       │           channels_out=10_226,
        │   │   │   │       │           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
        │   │   │   │       │   ),
        │   │   │   │       │   head_elements=SwiGLU.Config(
        │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │       channels_out=10,
        │   │   │   │       │       channels_hidden=256,
        │   │   │   │       │       expansion=4.0,
        │   │   │   │       │       round_to=256,
        │   │   │   │       │       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=1_024,
        │   │   │   │       │           channels_out=256,
        │   │   │   │       │           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=256,
        │   │   │   │       │           channels_out=10,
        │   │   │   │       │           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
        │   │   │   │       │   )
        │   │   │   │       ),
        │   │   │   │       dtype_autocast=torch.bfloat16
        │   │   │   │   ),
        │   │   │   │   BankCheckpoint.Config(
        │   │   │   │       base_dir=PosixPath('/opt/scratch/datasets/casmi2026'),
        │   │   │   │       path=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/ahmedberatozer/casmi26-v4b-models/fe_D.pt'),
        │   │   │   │       sha256='8ae5bee2c978f4f01ea7cfe29b514cd3e9112006ab51aa8b13f20e26e9365362',
        │   │   │   │       model=DreamsFP.Config(
        │   │   │   │       │   channels_out=10_226,
        │   │   │   │       │   num_elements=10,
        │   │   │   │       │   max_peaks=60,
        │   │   │   │       │   negative_mass_shift=2.014552934,
        │   │   │   │       │   precursor_intensity=1.1,
        │   │   │   │       │   backbone=DreamBackbone.Config(
        │   │   │   │       │   │   channels_fourier=980,
        │   │   │   │       │   │   channels_peak=44,
        │   │   │   │       │   │   max_mass=1000.0,
        │   │   │   │       │   │   fourier=Sequential.Config(
        │   │   │   │       │   │   │   channels_in=11_994,
        │   │   │   │       │   │   │   channels_out=980,
        │   │   │   │       │   │   │   elements=[
        │   │   │   │       │   │   │   │   Linear.Config(
        │   │   │   │       │   │   │   │       channels_in=11_994,
        │   │   │   │       │   │   │   │       channels_out=512,
        │   │   │   │       │   │   │   │       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
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   Dropout.Config(
        │   │   │   │       │   │   │   │       channels_in=512,
        │   │   │   │       │   │   │   │       channels_out=512,
        │   │   │   │       │   │   │   │       p=0.1,
        │   │   │   │       │   │   │   │       inplace=False
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   Activation.Config(
        │   │   │   │       │   │   │   │       channels_in=512,
        │   │   │   │       │   │   │   │       channels_out=512,
        │   │   │   │       │   │   │   │       act=priml.model.swiglu.relu
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   Linear.Config(
        │   │   │   │       │   │   │   │       channels_in=512,
        │   │   │   │       │   │   │   │       channels_out=512,
        │   │   │   │       │   │   │   │       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
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   Dropout.Config(
        │   │   │   │       │   │   │   │       channels_in=512,
        │   │   │   │       │   │   │   │       channels_out=512,
        │   │   │   │       │   │   │   │       p=0.1,
        │   │   │   │       │   │   │   │       inplace=False
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   Activation.Config(
        │   │   │   │       │   │   │   │       channels_in=512,
        │   │   │   │       │   │   │   │       channels_out=512,
        │   │   │   │       │   │   │   │       act=priml.model.swiglu.relu
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   Linear.Config(
        │   │   │   │       │   │   │   │       channels_in=512,
        │   │   │   │       │   │   │   │       channels_out=512,
        │   │   │   │       │   │   │   │       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
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   Dropout.Config(
        │   │   │   │       │   │   │   │       channels_in=512,
        │   │   │   │       │   │   │   │       channels_out=512,
        │   │   │   │       │   │   │   │       p=0.1,
        │   │   │   │       │   │   │   │       inplace=False
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   Activation.Config(
        │   │   │   │       │   │   │   │       channels_in=512,
        │   │   │   │       │   │   │   │       channels_out=512,
        │   │   │   │       │   │   │   │       act=priml.model.swiglu.relu
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   Linear.Config(
        │   │   │   │       │   │   │   │       channels_in=512,
        │   │   │   │       │   │   │   │       channels_out=512,
        │   │   │   │       │   │   │   │       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
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   Dropout.Config(
        │   │   │   │       │   │   │   │       channels_in=512,
        │   │   │   │       │   │   │   │       channels_out=512,
        │   │   │   │       │   │   │   │       p=0.1,
        │   │   │   │       │   │   │   │       inplace=False
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   Activation.Config(
        │   │   │   │       │   │   │   │       channels_in=512,
        │   │   │   │       │   │   │   │       channels_out=512,
        │   │   │   │       │   │   │   │       act=priml.model.swiglu.relu
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   Linear.Config(
        │   │   │   │       │   │   │   │       channels_in=512,
        │   │   │   │       │   │   │   │       channels_out=980,
        │   │   │   │       │   │   │   │       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
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   Activation.Config(
        │   │   │   │       │   │   │   │       channels_in=980,
        │   │   │   │       │   │   │   │       channels_out=980,
        │   │   │   │       │   │   │   │       act=priml.model.swiglu.relu
        │   │   │   │       │   │   │   │   )
        │   │   │   │       │   │   │   ],
        │   │   │   │       │   │   │   repeat=1,
        │   │   │   │       │   │   │   depth_index=()
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   peak=Sequential.Config(
        │   │   │   │       │   │       channels_in=2,
        │   │   │   │       │   │       channels_out=44,
        │   │   │   │       │   │       elements=[
        │   │   │   │       │   │           Linear.Config(
        │   │   │   │       │   │               channels_in=2,
        │   │   │   │       │   │               channels_out=44,
        │   │   │   │       │   │               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
        │   │   │   │       │   │           ),
        │   │   │   │       │   │           Activation.Config(
        │   │   │   │       │   │               channels_in=44,
        │   │   │   │       │   │               channels_out=44,
        │   │   │   │       │   │               act=priml.model.swiglu.relu
        │   │   │   │       │   │           )
        │   │   │   │       │   │       ],
        │   │   │   │       │   │       repeat=1,
        │   │   │   │       │   │       depth_index=()
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   basis=FourierFeatures.Config(
        │   │   │   │       │   │       num_fractional=4_999,
        │   │   │   │       │   │       num_integer=998,
        │   │   │   │       │   │       resolution=0.0001,
        │   │   │   │       │   │       channels_out=11_994
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   encoder=DreamEncoder.Config(
        │   │   │   │       │   │   │   channels_in=1_024,
        │   │   │   │       │   │   │   channels_out=1_024,
        │   │   │   │       │   │   │   num_layers=7,
        │   │   │   │       │   │   │   proj_in=Dropout.Config(
        │   │   │   │       │   │   │       channels_in=1_024,
        │   │   │   │       │   │   │       channels_out=1_024,
        │   │   │   │       │   │   │       p=0.1,
        │   │   │   │       │   │   │       inplace=False
        │   │   │   │       │   │   │   ),
        │   │   │   │       │   │   │   block=[
        │   │   │   │       │   │   │   │   TransformerBlock.Config(
        │   │   │   │       │   │   │   │       channels_in=1_024,
        │   │   │   │       │   │   │   │       channels_out=1_024,
        │   │   │   │       │   │   │   │       attn=DreamAttention.Config(
        │   │   │   │       │   │   │   │       │   channels_in=1_024,
        │   │   │   │       │   │   │   │       │   channels_out=1_024,
        │   │   │   │       │   │   │   │       │   attention=Attention.Config(
        │   │   │   │       │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │       channels_out=1_024,
        │   │   │   │       │   │   │   │       │       num_heads=8,
        │   │   │   │       │   │   │   │       │       channels_head=128,
        │   │   │   │       │   │   │   │       │       num_heads_kv=8,
        │   │   │   │       │   │   │   │       │       bias=False,
        │   │   │   │       │   │   │   │       │       dropout=0.1,
        │   │   │   │       │   │   │   │       │       causal=False,
        │   │   │   │       │   │   │   │       │       rope=None,
        │   │   │   │       │   │   │   │       │       rotate=None,
        │   │   │   │       │   │   │   │       │       norm_qk=None,
        │   │   │   │       │   │   │   │       │       share_qk_norm=True,
        │   │   │   │       │   │   │   │       │       norm_out=None,
        │   │   │   │       │   │   │   │       │       split_qkv_projection=False,
        │   │   │   │       │   │   │   │       │       proj_qkv=Linear.Config(
        │   │   │   │       │   │   │   │       │           channels_in=1_024,
        │   │   │   │       │   │   │   │       │           channels_out=3_072,
        │   │   │   │       │   │   │   │       │           bias=False,
        │   │   │   │       │   │   │   │       │           depth_index=(),
        │   │   │   │       │   │   │   │       │           device=None,
        │   │   │   │       │   │   │   │       │           dtype=None,
        │   │   │   │       │   │   │   │       │           init_weight=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │           shard=None
        │   │   │   │       │   │   │   │       │       ),
        │   │   │   │       │   │   │   │       │       proj_out=Linear.Config(
        │   │   │   │       │   │   │   │       │           channels_in=1_024,
        │   │   │   │       │   │   │   │       │           channels_out=1_024,
        │   │   │   │       │   │   │   │       │           bias=False,
        │   │   │   │       │   │   │   │       │           depth_index=(),
        │   │   │   │       │   │   │   │       │           device=None,
        │   │   │   │       │   │   │   │       │           dtype=None,
        │   │   │   │       │   │   │   │       │           init_weight=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │           shard=None
        │   │   │   │       │   │   │   │       │       ),
        │   │   │   │       │   │   │   │       │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │   │       │       depth_index=(),
        │   │   │   │       │   │   │   │       │       attn_kernel=DreamKernel.Config(
        │   │   │   │       │   │   │   │       │           dtype=torch.float32,
        │   │   │   │       │   │   │   │       │           kernel=SdpaFused.Config()
        │   │   │   │       │   │   │   │       │       )
        │   │   │   │       │   │   │   │       │   ),
        │   │   │   │       │   │   │   │       │   dropout=Dropout.Config(
        │   │   │   │       │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │       channels_out=1_024,
        │   │   │   │       │   │   │   │       │       p=0.1,
        │   │   │   │       │   │   │   │       │       inplace=False
        │   │   │   │       │   │   │   │       │   )
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       ffn=Sequential.Config(
        │   │   │   │       │   │   │   │       │   channels_in=1_024,
        │   │   │   │       │   │   │   │       │   channels_out=1_024,
        │   │   │   │       │   │   │   │       │   elements=[
        │   │   │   │       │   │   │   │       │   │   DreamFFN.Config(
        │   │   │   │       │   │   │   │       │   │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │       ffn=SwiGLU.Config(
        │   │   │   │       │   │   │   │       │   │           channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │           channels_out=1_024,
        │   │   │   │       │   │   │   │       │   │           channels_hidden=4_096,
        │   │   │   │       │   │   │   │       │   │           expansion=4.0,
        │   │   │   │       │   │   │   │       │   │           round_to=1,
        │   │   │   │       │   │   │   │       │   │           gate=False,
        │   │   │   │       │   │   │   │       │   │           bias=False,
        │   │   │   │       │   │   │   │       │   │           split_gate_projection=False,
        │   │   │   │       │   │   │   │       │   │           norm=None,
        │   │   │   │       │   │   │   │       │   │           act=priml.model.swiglu.relu,
        │   │   │   │       │   │   │   │       │   │           depth_index=(),
        │   │   │   │       │   │   │   │       │   │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │       │   │   │   │       │   │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │   │           up_proj=Linear.Config(
        │   │   │   │       │   │   │   │       │   │               channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │               channels_out=4_096,
        │   │   │   │       │   │   │   │       │   │               bias=False,
        │   │   │   │       │   │   │   │       │   │               depth_index=(),
        │   │   │   │       │   │   │   │       │   │               device=None,
        │   │   │   │       │   │   │   │       │   │               dtype=None,
        │   │   │   │       │   │   │   │       │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │   │       │   │               init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │   │               linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │   │               shard=None
        │   │   │   │       │   │   │   │       │   │           ),
        │   │   │   │       │   │   │   │       │   │           middle=Dropout.Config(
        │   │   │   │       │   │   │   │       │   │               channels_in=4_096,
        │   │   │   │       │   │   │   │       │   │               channels_out=4_096,
        │   │   │   │       │   │   │   │       │   │               p=0.1,
        │   │   │   │       │   │   │   │       │   │               inplace=False
        │   │   │   │       │   │   │   │       │   │           ),
        │   │   │   │       │   │   │   │       │   │           down_proj=Linear.Config(
        │   │   │   │       │   │   │   │       │   │               channels_in=4_096,
        │   │   │   │       │   │   │   │       │   │               channels_out=1_024,
        │   │   │   │       │   │   │   │       │   │               bias=False,
        │   │   │   │       │   │   │   │       │   │               depth_index=(),
        │   │   │   │       │   │   │   │       │   │               device=None,
        │   │   │   │       │   │   │   │       │   │               dtype=None,
        │   │   │   │       │   │   │   │       │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │   │       │   │               init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │   │               linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │   │               shard=None
        │   │   │   │       │   │   │   │       │   │           ),
        │   │   │   │       │   │   │   │       │   │           shard=None
        │   │   │   │       │   │   │   │       │   │       ),
        │   │   │   │       │   │   │   │       │   │       init_final=WidthNormal(width=1_024)
        │   │   │   │       │   │   │   │       │   │   ),
        │   │   │   │       │   │   │   │       │   │   Dropout.Config(
        │   │   │   │       │   │   │   │       │   │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │       channels_out=1_024,
        │   │   │   │       │   │   │   │       │   │       p=0.1,
        │   │   │   │       │   │   │   │       │   │       inplace=False
        │   │   │   │       │   │   │   │       │   │   )
        │   │   │   │       │   │   │   │       │   ],
        │   │   │   │       │   │   │   │       │   repeat=1,
        │   │   │   │       │   │   │   │       │   depth_index=((0, 7),)
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       norm1=LayerNorm.Config(
        │   │   │   │       │   │   │   │           channels_in=1_024,
        │   │   │   │       │   │   │   │           channels_out=1_024,
        │   │   │   │       │   │   │   │           eps=1e-05,
        │   │   │   │       │   │   │   │           elementwise_affine=True,
        │   │   │   │       │   │   │   │           device=None,
        │   │   │   │       │   │   │   │           dtype=None
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       norm2=LayerNorm.Config(
        │   │   │   │       │   │   │   │           channels_in=1_024,
        │   │   │   │       │   │   │   │           channels_out=1_024,
        │   │   │   │       │   │   │   │           eps=1e-05,
        │   │   │   │       │   │   │   │           elementwise_affine=True,
        │   │   │   │       │   │   │   │           device=None,
        │   │   │   │       │   │   │   │           dtype=None
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       prenorm=True,
        │   │   │   │       │   │   │   │       checkpoint=False,
        │   │   │   │       │   │   │   │       depth_index=((0, 7),)
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   TransformerBlock.Config(
        │   │   │   │       │   │   │   │       channels_in=1_024,
        │   │   │   │       │   │   │   │       channels_out=1_024,
        │   │   │   │       │   │   │   │       attn=DreamAttention.Config(
        │   │   │   │       │   │   │   │       │   channels_in=1_024,
        │   │   │   │       │   │   │   │       │   channels_out=1_024,
        │   │   │   │       │   │   │   │       │   attention=Attention.Config(
        │   │   │   │       │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │       channels_out=1_024,
        │   │   │   │       │   │   │   │       │       num_heads=8,
        │   │   │   │       │   │   │   │       │       channels_head=128,
        │   │   │   │       │   │   │   │       │       num_heads_kv=8,
        │   │   │   │       │   │   │   │       │       bias=False,
        │   │   │   │       │   │   │   │       │       dropout=0.1,
        │   │   │   │       │   │   │   │       │       causal=False,
        │   │   │   │       │   │   │   │       │       rope=None,
        │   │   │   │       │   │   │   │       │       rotate=None,
        │   │   │   │       │   │   │   │       │       norm_qk=None,
        │   │   │   │       │   │   │   │       │       share_qk_norm=True,
        │   │   │   │       │   │   │   │       │       norm_out=None,
        │   │   │   │       │   │   │   │       │       split_qkv_projection=False,
        │   │   │   │       │   │   │   │       │       proj_qkv=Linear.Config(
        │   │   │   │       │   │   │   │       │           channels_in=1_024,
        │   │   │   │       │   │   │   │       │           channels_out=3_072,
        │   │   │   │       │   │   │   │       │           bias=False,
        │   │   │   │       │   │   │   │       │           depth_index=(),
        │   │   │   │       │   │   │   │       │           device=None,
        │   │   │   │       │   │   │   │       │           dtype=None,
        │   │   │   │       │   │   │   │       │           init_weight=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │           shard=None
        │   │   │   │       │   │   │   │       │       ),
        │   │   │   │       │   │   │   │       │       proj_out=Linear.Config(
        │   │   │   │       │   │   │   │       │           channels_in=1_024,
        │   │   │   │       │   │   │   │       │           channels_out=1_024,
        │   │   │   │       │   │   │   │       │           bias=False,
        │   │   │   │       │   │   │   │       │           depth_index=(),
        │   │   │   │       │   │   │   │       │           device=None,
        │   │   │   │       │   │   │   │       │           dtype=None,
        │   │   │   │       │   │   │   │       │           init_weight=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │           shard=None
        │   │   │   │       │   │   │   │       │       ),
        │   │   │   │       │   │   │   │       │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │   │       │       depth_index=(),
        │   │   │   │       │   │   │   │       │       attn_kernel=DreamKernel.Config(
        │   │   │   │       │   │   │   │       │           dtype=torch.float32,
        │   │   │   │       │   │   │   │       │           kernel=SdpaFused.Config()
        │   │   │   │       │   │   │   │       │       )
        │   │   │   │       │   │   │   │       │   ),
        │   │   │   │       │   │   │   │       │   dropout=Dropout.Config(
        │   │   │   │       │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │       channels_out=1_024,
        │   │   │   │       │   │   │   │       │       p=0.1,
        │   │   │   │       │   │   │   │       │       inplace=False
        │   │   │   │       │   │   │   │       │   )
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       ffn=Sequential.Config(
        │   │   │   │       │   │   │   │       │   channels_in=1_024,
        │   │   │   │       │   │   │   │       │   channels_out=1_024,
        │   │   │   │       │   │   │   │       │   elements=[
        │   │   │   │       │   │   │   │       │   │   DreamFFN.Config(
        │   │   │   │       │   │   │   │       │   │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │       ffn=SwiGLU.Config(
        │   │   │   │       │   │   │   │       │   │           channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │           channels_out=1_024,
        │   │   │   │       │   │   │   │       │   │           channels_hidden=4_096,
        │   │   │   │       │   │   │   │       │   │           expansion=4.0,
        │   │   │   │       │   │   │   │       │   │           round_to=1,
        │   │   │   │       │   │   │   │       │   │           gate=False,
        │   │   │   │       │   │   │   │       │   │           bias=False,
        │   │   │   │       │   │   │   │       │   │           split_gate_projection=False,
        │   │   │   │       │   │   │   │       │   │           norm=None,
        │   │   │   │       │   │   │   │       │   │           act=priml.model.swiglu.relu,
        │   │   │   │       │   │   │   │       │   │           depth_index=(),
        │   │   │   │       │   │   │   │       │   │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │       │   │   │   │       │   │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │   │           up_proj=Linear.Config(
        │   │   │   │       │   │   │   │       │   │               channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │               channels_out=4_096,
        │   │   │   │       │   │   │   │       │   │               bias=False,
        │   │   │   │       │   │   │   │       │   │               depth_index=(),
        │   │   │   │       │   │   │   │       │   │               device=None,
        │   │   │   │       │   │   │   │       │   │               dtype=None,
        │   │   │   │       │   │   │   │       │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │   │       │   │               init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │   │               linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │   │               shard=None
        │   │   │   │       │   │   │   │       │   │           ),
        │   │   │   │       │   │   │   │       │   │           middle=Dropout.Config(
        │   │   │   │       │   │   │   │       │   │               channels_in=4_096,
        │   │   │   │       │   │   │   │       │   │               channels_out=4_096,
        │   │   │   │       │   │   │   │       │   │               p=0.1,
        │   │   │   │       │   │   │   │       │   │               inplace=False
        │   │   │   │       │   │   │   │       │   │           ),
        │   │   │   │       │   │   │   │       │   │           down_proj=Linear.Config(
        │   │   │   │       │   │   │   │       │   │               channels_in=4_096,
        │   │   │   │       │   │   │   │       │   │               channels_out=1_024,
        │   │   │   │       │   │   │   │       │   │               bias=False,
        │   │   │   │       │   │   │   │       │   │               depth_index=(),
        │   │   │   │       │   │   │   │       │   │               device=None,
        │   │   │   │       │   │   │   │       │   │               dtype=None,
        │   │   │   │       │   │   │   │       │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │   │       │   │               init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │   │               linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │   │               shard=None
        │   │   │   │       │   │   │   │       │   │           ),
        │   │   │   │       │   │   │   │       │   │           shard=None
        │   │   │   │       │   │   │   │       │   │       ),
        │   │   │   │       │   │   │   │       │   │       init_final=WidthNormal(width=1_024)
        │   │   │   │       │   │   │   │       │   │   ),
        │   │   │   │       │   │   │   │       │   │   Dropout.Config(
        │   │   │   │       │   │   │   │       │   │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │       channels_out=1_024,
        │   │   │   │       │   │   │   │       │   │       p=0.1,
        │   │   │   │       │   │   │   │       │   │       inplace=False
        │   │   │   │       │   │   │   │       │   │   )
        │   │   │   │       │   │   │   │       │   ],
        │   │   │   │       │   │   │   │       │   repeat=1,
        │   │   │   │       │   │   │   │       │   depth_index=((1, 7),)
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       norm1=LayerNorm.Config(
        │   │   │   │       │   │   │   │           channels_in=1_024,
        │   │   │   │       │   │   │   │           channels_out=1_024,
        │   │   │   │       │   │   │   │           eps=1e-05,
        │   │   │   │       │   │   │   │           elementwise_affine=True,
        │   │   │   │       │   │   │   │           device=None,
        │   │   │   │       │   │   │   │           dtype=None
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       norm2=LayerNorm.Config(
        │   │   │   │       │   │   │   │           channels_in=1_024,
        │   │   │   │       │   │   │   │           channels_out=1_024,
        │   │   │   │       │   │   │   │           eps=1e-05,
        │   │   │   │       │   │   │   │           elementwise_affine=True,
        │   │   │   │       │   │   │   │           device=None,
        │   │   │   │       │   │   │   │           dtype=None
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       prenorm=True,
        │   │   │   │       │   │   │   │       checkpoint=False,
        │   │   │   │       │   │   │   │       depth_index=((1, 7),)
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   TransformerBlock.Config(
        │   │   │   │       │   │   │   │       channels_in=1_024,
        │   │   │   │       │   │   │   │       channels_out=1_024,
        │   │   │   │       │   │   │   │       attn=DreamAttention.Config(
        │   │   │   │       │   │   │   │       │   channels_in=1_024,
        │   │   │   │       │   │   │   │       │   channels_out=1_024,
        │   │   │   │       │   │   │   │       │   attention=Attention.Config(
        │   │   │   │       │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │       channels_out=1_024,
        │   │   │   │       │   │   │   │       │       num_heads=8,
        │   │   │   │       │   │   │   │       │       channels_head=128,
        │   │   │   │       │   │   │   │       │       num_heads_kv=8,
        │   │   │   │       │   │   │   │       │       bias=False,
        │   │   │   │       │   │   │   │       │       dropout=0.1,
        │   │   │   │       │   │   │   │       │       causal=False,
        │   │   │   │       │   │   │   │       │       rope=None,
        │   │   │   │       │   │   │   │       │       rotate=None,
        │   │   │   │       │   │   │   │       │       norm_qk=None,
        │   │   │   │       │   │   │   │       │       share_qk_norm=True,
        │   │   │   │       │   │   │   │       │       norm_out=None,
        │   │   │   │       │   │   │   │       │       split_qkv_projection=False,
        │   │   │   │       │   │   │   │       │       proj_qkv=Linear.Config(
        │   │   │   │       │   │   │   │       │           channels_in=1_024,
        │   │   │   │       │   │   │   │       │           channels_out=3_072,
        │   │   │   │       │   │   │   │       │           bias=False,
        │   │   │   │       │   │   │   │       │           depth_index=(),
        │   │   │   │       │   │   │   │       │           device=None,
        │   │   │   │       │   │   │   │       │           dtype=None,
        │   │   │   │       │   │   │   │       │           init_weight=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │           shard=None
        │   │   │   │       │   │   │   │       │       ),
        │   │   │   │       │   │   │   │       │       proj_out=Linear.Config(
        │   │   │   │       │   │   │   │       │           channels_in=1_024,
        │   │   │   │       │   │   │   │       │           channels_out=1_024,
        │   │   │   │       │   │   │   │       │           bias=False,
        │   │   │   │       │   │   │   │       │           depth_index=(),
        │   │   │   │       │   │   │   │       │           device=None,
        │   │   │   │       │   │   │   │       │           dtype=None,
        │   │   │   │       │   │   │   │       │           init_weight=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │           shard=None
        │   │   │   │       │   │   │   │       │       ),
        │   │   │   │       │   │   │   │       │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │   │       │       depth_index=(),
        │   │   │   │       │   │   │   │       │       attn_kernel=DreamKernel.Config(
        │   │   │   │       │   │   │   │       │           dtype=torch.float32,
        │   │   │   │       │   │   │   │       │           kernel=SdpaFused.Config()
        │   │   │   │       │   │   │   │       │       )
        │   │   │   │       │   │   │   │       │   ),
        │   │   │   │       │   │   │   │       │   dropout=Dropout.Config(
        │   │   │   │       │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │       channels_out=1_024,
        │   │   │   │       │   │   │   │       │       p=0.1,
        │   │   │   │       │   │   │   │       │       inplace=False
        │   │   │   │       │   │   │   │       │   )
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       ffn=Sequential.Config(
        │   │   │   │       │   │   │   │       │   channels_in=1_024,
        │   │   │   │       │   │   │   │       │   channels_out=1_024,
        │   │   │   │       │   │   │   │       │   elements=[
        │   │   │   │       │   │   │   │       │   │   DreamFFN.Config(
        │   │   │   │       │   │   │   │       │   │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │       ffn=SwiGLU.Config(
        │   │   │   │       │   │   │   │       │   │           channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │           channels_out=1_024,
        │   │   │   │       │   │   │   │       │   │           channels_hidden=4_096,
        │   │   │   │       │   │   │   │       │   │           expansion=4.0,
        │   │   │   │       │   │   │   │       │   │           round_to=1,
        │   │   │   │       │   │   │   │       │   │           gate=False,
        │   │   │   │       │   │   │   │       │   │           bias=False,
        │   │   │   │       │   │   │   │       │   │           split_gate_projection=False,
        │   │   │   │       │   │   │   │       │   │           norm=None,
        │   │   │   │       │   │   │   │       │   │           act=priml.model.swiglu.relu,
        │   │   │   │       │   │   │   │       │   │           depth_index=(),
        │   │   │   │       │   │   │   │       │   │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │       │   │   │   │       │   │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │   │           up_proj=Linear.Config(
        │   │   │   │       │   │   │   │       │   │               channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │               channels_out=4_096,
        │   │   │   │       │   │   │   │       │   │               bias=False,
        │   │   │   │       │   │   │   │       │   │               depth_index=(),
        │   │   │   │       │   │   │   │       │   │               device=None,
        │   │   │   │       │   │   │   │       │   │               dtype=None,
        │   │   │   │       │   │   │   │       │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │   │       │   │               init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │   │               linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │   │               shard=None
        │   │   │   │       │   │   │   │       │   │           ),
        │   │   │   │       │   │   │   │       │   │           middle=Dropout.Config(
        │   │   │   │       │   │   │   │       │   │               channels_in=4_096,
        │   │   │   │       │   │   │   │       │   │               channels_out=4_096,
        │   │   │   │       │   │   │   │       │   │               p=0.1,
        │   │   │   │       │   │   │   │       │   │               inplace=False
        │   │   │   │       │   │   │   │       │   │           ),
        │   │   │   │       │   │   │   │       │   │           down_proj=Linear.Config(
        │   │   │   │       │   │   │   │       │   │               channels_in=4_096,
        │   │   │   │       │   │   │   │       │   │               channels_out=1_024,
        │   │   │   │       │   │   │   │       │   │               bias=False,
        │   │   │   │       │   │   │   │       │   │               depth_index=(),
        │   │   │   │       │   │   │   │       │   │               device=None,
        │   │   │   │       │   │   │   │       │   │               dtype=None,
        │   │   │   │       │   │   │   │       │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │   │       │   │               init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │   │               linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │   │               shard=None
        │   │   │   │       │   │   │   │       │   │           ),
        │   │   │   │       │   │   │   │       │   │           shard=None
        │   │   │   │       │   │   │   │       │   │       ),
        │   │   │   │       │   │   │   │       │   │       init_final=WidthNormal(width=1_024)
        │   │   │   │       │   │   │   │       │   │   ),
        │   │   │   │       │   │   │   │       │   │   Dropout.Config(
        │   │   │   │       │   │   │   │       │   │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │       channels_out=1_024,
        │   │   │   │       │   │   │   │       │   │       p=0.1,
        │   │   │   │       │   │   │   │       │   │       inplace=False
        │   │   │   │       │   │   │   │       │   │   )
        │   │   │   │       │   │   │   │       │   ],
        │   │   │   │       │   │   │   │       │   repeat=1,
        │   │   │   │       │   │   │   │       │   depth_index=((2, 7),)
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       norm1=LayerNorm.Config(
        │   │   │   │       │   │   │   │           channels_in=1_024,
        │   │   │   │       │   │   │   │           channels_out=1_024,
        │   │   │   │       │   │   │   │           eps=1e-05,
        │   │   │   │       │   │   │   │           elementwise_affine=True,
        │   │   │   │       │   │   │   │           device=None,
        │   │   │   │       │   │   │   │           dtype=None
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       norm2=LayerNorm.Config(
        │   │   │   │       │   │   │   │           channels_in=1_024,
        │   │   │   │       │   │   │   │           channels_out=1_024,
        │   │   │   │       │   │   │   │           eps=1e-05,
        │   │   │   │       │   │   │   │           elementwise_affine=True,
        │   │   │   │       │   │   │   │           device=None,
        │   │   │   │       │   │   │   │           dtype=None
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       prenorm=True,
        │   │   │   │       │   │   │   │       checkpoint=False,
        │   │   │   │       │   │   │   │       depth_index=((2, 7),)
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   TransformerBlock.Config(
        │   │   │   │       │   │   │   │       channels_in=1_024,
        │   │   │   │       │   │   │   │       channels_out=1_024,
        │   │   │   │       │   │   │   │       attn=DreamAttention.Config(
        │   │   │   │       │   │   │   │       │   channels_in=1_024,
        │   │   │   │       │   │   │   │       │   channels_out=1_024,
        │   │   │   │       │   │   │   │       │   attention=Attention.Config(
        │   │   │   │       │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │       channels_out=1_024,
        │   │   │   │       │   │   │   │       │       num_heads=8,
        │   │   │   │       │   │   │   │       │       channels_head=128,
        │   │   │   │       │   │   │   │       │       num_heads_kv=8,
        │   │   │   │       │   │   │   │       │       bias=False,
        │   │   │   │       │   │   │   │       │       dropout=0.1,
        │   │   │   │       │   │   │   │       │       causal=False,
        │   │   │   │       │   │   │   │       │       rope=None,
        │   │   │   │       │   │   │   │       │       rotate=None,
        │   │   │   │       │   │   │   │       │       norm_qk=None,
        │   │   │   │       │   │   │   │       │       share_qk_norm=True,
        │   │   │   │       │   │   │   │       │       norm_out=None,
        │   │   │   │       │   │   │   │       │       split_qkv_projection=False,
        │   │   │   │       │   │   │   │       │       proj_qkv=Linear.Config(
        │   │   │   │       │   │   │   │       │           channels_in=1_024,
        │   │   │   │       │   │   │   │       │           channels_out=3_072,
        │   │   │   │       │   │   │   │       │           bias=False,
        │   │   │   │       │   │   │   │       │           depth_index=(),
        │   │   │   │       │   │   │   │       │           device=None,
        │   │   │   │       │   │   │   │       │           dtype=None,
        │   │   │   │       │   │   │   │       │           init_weight=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │           shard=None
        │   │   │   │       │   │   │   │       │       ),
        │   │   │   │       │   │   │   │       │       proj_out=Linear.Config(
        │   │   │   │       │   │   │   │       │           channels_in=1_024,
        │   │   │   │       │   │   │   │       │           channels_out=1_024,
        │   │   │   │       │   │   │   │       │           bias=False,
        │   │   │   │       │   │   │   │       │           depth_index=(),
        │   │   │   │       │   │   │   │       │           device=None,
        │   │   │   │       │   │   │   │       │           dtype=None,
        │   │   │   │       │   │   │   │       │           init_weight=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │           shard=None
        │   │   │   │       │   │   │   │       │       ),
        │   │   │   │       │   │   │   │       │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │   │       │       depth_index=(),
        │   │   │   │       │   │   │   │       │       attn_kernel=DreamKernel.Config(
        │   │   │   │       │   │   │   │       │           dtype=torch.float32,
        │   │   │   │       │   │   │   │       │           kernel=SdpaFused.Config()
        │   │   │   │       │   │   │   │       │       )
        │   │   │   │       │   │   │   │       │   ),
        │   │   │   │       │   │   │   │       │   dropout=Dropout.Config(
        │   │   │   │       │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │       channels_out=1_024,
        │   │   │   │       │   │   │   │       │       p=0.1,
        │   │   │   │       │   │   │   │       │       inplace=False
        │   │   │   │       │   │   │   │       │   )
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       ffn=Sequential.Config(
        │   │   │   │       │   │   │   │       │   channels_in=1_024,
        │   │   │   │       │   │   │   │       │   channels_out=1_024,
        │   │   │   │       │   │   │   │       │   elements=[
        │   │   │   │       │   │   │   │       │   │   DreamFFN.Config(
        │   │   │   │       │   │   │   │       │   │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │       ffn=SwiGLU.Config(
        │   │   │   │       │   │   │   │       │   │           channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │           channels_out=1_024,
        │   │   │   │       │   │   │   │       │   │           channels_hidden=4_096,
        │   │   │   │       │   │   │   │       │   │           expansion=4.0,
        │   │   │   │       │   │   │   │       │   │           round_to=1,
        │   │   │   │       │   │   │   │       │   │           gate=False,
        │   │   │   │       │   │   │   │       │   │           bias=False,
        │   │   │   │       │   │   │   │       │   │           split_gate_projection=False,
        │   │   │   │       │   │   │   │       │   │           norm=None,
        │   │   │   │       │   │   │   │       │   │           act=priml.model.swiglu.relu,
        │   │   │   │       │   │   │   │       │   │           depth_index=(),
        │   │   │   │       │   │   │   │       │   │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │       │   │   │   │       │   │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │   │           up_proj=Linear.Config(
        │   │   │   │       │   │   │   │       │   │               channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │               channels_out=4_096,
        │   │   │   │       │   │   │   │       │   │               bias=False,
        │   │   │   │       │   │   │   │       │   │               depth_index=(),
        │   │   │   │       │   │   │   │       │   │               device=None,
        │   │   │   │       │   │   │   │       │   │               dtype=None,
        │   │   │   │       │   │   │   │       │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │   │       │   │               init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │   │               linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │   │               shard=None
        │   │   │   │       │   │   │   │       │   │           ),
        │   │   │   │       │   │   │   │       │   │           middle=Dropout.Config(
        │   │   │   │       │   │   │   │       │   │               channels_in=4_096,
        │   │   │   │       │   │   │   │       │   │               channels_out=4_096,
        │   │   │   │       │   │   │   │       │   │               p=0.1,
        │   │   │   │       │   │   │   │       │   │               inplace=False
        │   │   │   │       │   │   │   │       │   │           ),
        │   │   │   │       │   │   │   │       │   │           down_proj=Linear.Config(
        │   │   │   │       │   │   │   │       │   │               channels_in=4_096,
        │   │   │   │       │   │   │   │       │   │               channels_out=1_024,
        │   │   │   │       │   │   │   │       │   │               bias=False,
        │   │   │   │       │   │   │   │       │   │               depth_index=(),
        │   │   │   │       │   │   │   │       │   │               device=None,
        │   │   │   │       │   │   │   │       │   │               dtype=None,
        │   │   │   │       │   │   │   │       │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │   │       │   │               init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │   │               linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │   │               shard=None
        │   │   │   │       │   │   │   │       │   │           ),
        │   │   │   │       │   │   │   │       │   │           shard=None
        │   │   │   │       │   │   │   │       │   │       ),
        │   │   │   │       │   │   │   │       │   │       init_final=WidthNormal(width=1_024)
        │   │   │   │       │   │   │   │       │   │   ),
        │   │   │   │       │   │   │   │       │   │   Dropout.Config(
        │   │   │   │       │   │   │   │       │   │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │       channels_out=1_024,
        │   │   │   │       │   │   │   │       │   │       p=0.1,
        │   │   │   │       │   │   │   │       │   │       inplace=False
        │   │   │   │       │   │   │   │       │   │   )
        │   │   │   │       │   │   │   │       │   ],
        │   │   │   │       │   │   │   │       │   repeat=1,
        │   │   │   │       │   │   │   │       │   depth_index=((3, 7),)
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       norm1=LayerNorm.Config(
        │   │   │   │       │   │   │   │           channels_in=1_024,
        │   │   │   │       │   │   │   │           channels_out=1_024,
        │   │   │   │       │   │   │   │           eps=1e-05,
        │   │   │   │       │   │   │   │           elementwise_affine=True,
        │   │   │   │       │   │   │   │           device=None,
        │   │   │   │       │   │   │   │           dtype=None
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       norm2=LayerNorm.Config(
        │   │   │   │       │   │   │   │           channels_in=1_024,
        │   │   │   │       │   │   │   │           channels_out=1_024,
        │   │   │   │       │   │   │   │           eps=1e-05,
        │   │   │   │       │   │   │   │           elementwise_affine=True,
        │   │   │   │       │   │   │   │           device=None,
        │   │   │   │       │   │   │   │           dtype=None
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       prenorm=True,
        │   │   │   │       │   │   │   │       checkpoint=False,
        │   │   │   │       │   │   │   │       depth_index=((3, 7),)
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   TransformerBlock.Config(
        │   │   │   │       │   │   │   │       channels_in=1_024,
        │   │   │   │       │   │   │   │       channels_out=1_024,
        │   │   │   │       │   │   │   │       attn=DreamAttention.Config(
        │   │   │   │       │   │   │   │       │   channels_in=1_024,
        │   │   │   │       │   │   │   │       │   channels_out=1_024,
        │   │   │   │       │   │   │   │       │   attention=Attention.Config(
        │   │   │   │       │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │       channels_out=1_024,
        │   │   │   │       │   │   │   │       │       num_heads=8,
        │   │   │   │       │   │   │   │       │       channels_head=128,
        │   │   │   │       │   │   │   │       │       num_heads_kv=8,
        │   │   │   │       │   │   │   │       │       bias=False,
        │   │   │   │       │   │   │   │       │       dropout=0.1,
        │   │   │   │       │   │   │   │       │       causal=False,
        │   │   │   │       │   │   │   │       │       rope=None,
        │   │   │   │       │   │   │   │       │       rotate=None,
        │   │   │   │       │   │   │   │       │       norm_qk=None,
        │   │   │   │       │   │   │   │       │       share_qk_norm=True,
        │   │   │   │       │   │   │   │       │       norm_out=None,
        │   │   │   │       │   │   │   │       │       split_qkv_projection=False,
        │   │   │   │       │   │   │   │       │       proj_qkv=Linear.Config(
        │   │   │   │       │   │   │   │       │           channels_in=1_024,
        │   │   │   │       │   │   │   │       │           channels_out=3_072,
        │   │   │   │       │   │   │   │       │           bias=False,
        │   │   │   │       │   │   │   │       │           depth_index=(),
        │   │   │   │       │   │   │   │       │           device=None,
        │   │   │   │       │   │   │   │       │           dtype=None,
        │   │   │   │       │   │   │   │       │           init_weight=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │           shard=None
        │   │   │   │       │   │   │   │       │       ),
        │   │   │   │       │   │   │   │       │       proj_out=Linear.Config(
        │   │   │   │       │   │   │   │       │           channels_in=1_024,
        │   │   │   │       │   │   │   │       │           channels_out=1_024,
        │   │   │   │       │   │   │   │       │           bias=False,
        │   │   │   │       │   │   │   │       │           depth_index=(),
        │   │   │   │       │   │   │   │       │           device=None,
        │   │   │   │       │   │   │   │       │           dtype=None,
        │   │   │   │       │   │   │   │       │           init_weight=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │           shard=None
        │   │   │   │       │   │   │   │       │       ),
        │   │   │   │       │   │   │   │       │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │   │       │       depth_index=(),
        │   │   │   │       │   │   │   │       │       attn_kernel=DreamKernel.Config(
        │   │   │   │       │   │   │   │       │           dtype=torch.float32,
        │   │   │   │       │   │   │   │       │           kernel=SdpaFused.Config()
        │   │   │   │       │   │   │   │       │       )
        │   │   │   │       │   │   │   │       │   ),
        │   │   │   │       │   │   │   │       │   dropout=Dropout.Config(
        │   │   │   │       │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │       channels_out=1_024,
        │   │   │   │       │   │   │   │       │       p=0.1,
        │   │   │   │       │   │   │   │       │       inplace=False
        │   │   │   │       │   │   │   │       │   )
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       ffn=Sequential.Config(
        │   │   │   │       │   │   │   │       │   channels_in=1_024,
        │   │   │   │       │   │   │   │       │   channels_out=1_024,
        │   │   │   │       │   │   │   │       │   elements=[
        │   │   │   │       │   │   │   │       │   │   DreamFFN.Config(
        │   │   │   │       │   │   │   │       │   │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │       ffn=SwiGLU.Config(
        │   │   │   │       │   │   │   │       │   │           channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │           channels_out=1_024,
        │   │   │   │       │   │   │   │       │   │           channels_hidden=4_096,
        │   │   │   │       │   │   │   │       │   │           expansion=4.0,
        │   │   │   │       │   │   │   │       │   │           round_to=1,
        │   │   │   │       │   │   │   │       │   │           gate=False,
        │   │   │   │       │   │   │   │       │   │           bias=False,
        │   │   │   │       │   │   │   │       │   │           split_gate_projection=False,
        │   │   │   │       │   │   │   │       │   │           norm=None,
        │   │   │   │       │   │   │   │       │   │           act=priml.model.swiglu.relu,
        │   │   │   │       │   │   │   │       │   │           depth_index=(),
        │   │   │   │       │   │   │   │       │   │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │       │   │   │   │       │   │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │   │           up_proj=Linear.Config(
        │   │   │   │       │   │   │   │       │   │               channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │               channels_out=4_096,
        │   │   │   │       │   │   │   │       │   │               bias=False,
        │   │   │   │       │   │   │   │       │   │               depth_index=(),
        │   │   │   │       │   │   │   │       │   │               device=None,
        │   │   │   │       │   │   │   │       │   │               dtype=None,
        │   │   │   │       │   │   │   │       │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │   │       │   │               init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │   │               linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │   │               shard=None
        │   │   │   │       │   │   │   │       │   │           ),
        │   │   │   │       │   │   │   │       │   │           middle=Dropout.Config(
        │   │   │   │       │   │   │   │       │   │               channels_in=4_096,
        │   │   │   │       │   │   │   │       │   │               channels_out=4_096,
        │   │   │   │       │   │   │   │       │   │               p=0.1,
        │   │   │   │       │   │   │   │       │   │               inplace=False
        │   │   │   │       │   │   │   │       │   │           ),
        │   │   │   │       │   │   │   │       │   │           down_proj=Linear.Config(
        │   │   │   │       │   │   │   │       │   │               channels_in=4_096,
        │   │   │   │       │   │   │   │       │   │               channels_out=1_024,
        │   │   │   │       │   │   │   │       │   │               bias=False,
        │   │   │   │       │   │   │   │       │   │               depth_index=(),
        │   │   │   │       │   │   │   │       │   │               device=None,
        │   │   │   │       │   │   │   │       │   │               dtype=None,
        │   │   │   │       │   │   │   │       │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │   │       │   │               init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │   │               linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │   │               shard=None
        │   │   │   │       │   │   │   │       │   │           ),
        │   │   │   │       │   │   │   │       │   │           shard=None
        │   │   │   │       │   │   │   │       │   │       ),
        │   │   │   │       │   │   │   │       │   │       init_final=WidthNormal(width=1_024)
        │   │   │   │       │   │   │   │       │   │   ),
        │   │   │   │       │   │   │   │       │   │   Dropout.Config(
        │   │   │   │       │   │   │   │       │   │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │       channels_out=1_024,
        │   │   │   │       │   │   │   │       │   │       p=0.1,
        │   │   │   │       │   │   │   │       │   │       inplace=False
        │   │   │   │       │   │   │   │       │   │   )
        │   │   │   │       │   │   │   │       │   ],
        │   │   │   │       │   │   │   │       │   repeat=1,
        │   │   │   │       │   │   │   │       │   depth_index=((4, 7),)
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       norm1=LayerNorm.Config(
        │   │   │   │       │   │   │   │           channels_in=1_024,
        │   │   │   │       │   │   │   │           channels_out=1_024,
        │   │   │   │       │   │   │   │           eps=1e-05,
        │   │   │   │       │   │   │   │           elementwise_affine=True,
        │   │   │   │       │   │   │   │           device=None,
        │   │   │   │       │   │   │   │           dtype=None
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       norm2=LayerNorm.Config(
        │   │   │   │       │   │   │   │           channels_in=1_024,
        │   │   │   │       │   │   │   │           channels_out=1_024,
        │   │   │   │       │   │   │   │           eps=1e-05,
        │   │   │   │       │   │   │   │           elementwise_affine=True,
        │   │   │   │       │   │   │   │           device=None,
        │   │   │   │       │   │   │   │           dtype=None
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       prenorm=True,
        │   │   │   │       │   │   │   │       checkpoint=False,
        │   │   │   │       │   │   │   │       depth_index=((4, 7),)
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   TransformerBlock.Config(
        │   │   │   │       │   │   │   │       channels_in=1_024,
        │   │   │   │       │   │   │   │       channels_out=1_024,
        │   │   │   │       │   │   │   │       attn=DreamAttention.Config(
        │   │   │   │       │   │   │   │       │   channels_in=1_024,
        │   │   │   │       │   │   │   │       │   channels_out=1_024,
        │   │   │   │       │   │   │   │       │   attention=Attention.Config(
        │   │   │   │       │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │       channels_out=1_024,
        │   │   │   │       │   │   │   │       │       num_heads=8,
        │   │   │   │       │   │   │   │       │       channels_head=128,
        │   │   │   │       │   │   │   │       │       num_heads_kv=8,
        │   │   │   │       │   │   │   │       │       bias=False,
        │   │   │   │       │   │   │   │       │       dropout=0.1,
        │   │   │   │       │   │   │   │       │       causal=False,
        │   │   │   │       │   │   │   │       │       rope=None,
        │   │   │   │       │   │   │   │       │       rotate=None,
        │   │   │   │       │   │   │   │       │       norm_qk=None,
        │   │   │   │       │   │   │   │       │       share_qk_norm=True,
        │   │   │   │       │   │   │   │       │       norm_out=None,
        │   │   │   │       │   │   │   │       │       split_qkv_projection=False,
        │   │   │   │       │   │   │   │       │       proj_qkv=Linear.Config(
        │   │   │   │       │   │   │   │       │           channels_in=1_024,
        │   │   │   │       │   │   │   │       │           channels_out=3_072,
        │   │   │   │       │   │   │   │       │           bias=False,
        │   │   │   │       │   │   │   │       │           depth_index=(),
        │   │   │   │       │   │   │   │       │           device=None,
        │   │   │   │       │   │   │   │       │           dtype=None,
        │   │   │   │       │   │   │   │       │           init_weight=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │           shard=None
        │   │   │   │       │   │   │   │       │       ),
        │   │   │   │       │   │   │   │       │       proj_out=Linear.Config(
        │   │   │   │       │   │   │   │       │           channels_in=1_024,
        │   │   │   │       │   │   │   │       │           channels_out=1_024,
        │   │   │   │       │   │   │   │       │           bias=False,
        │   │   │   │       │   │   │   │       │           depth_index=(),
        │   │   │   │       │   │   │   │       │           device=None,
        │   │   │   │       │   │   │   │       │           dtype=None,
        │   │   │   │       │   │   │   │       │           init_weight=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │           shard=None
        │   │   │   │       │   │   │   │       │       ),
        │   │   │   │       │   │   │   │       │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │   │       │       depth_index=(),
        │   │   │   │       │   │   │   │       │       attn_kernel=DreamKernel.Config(
        │   │   │   │       │   │   │   │       │           dtype=torch.float32,
        │   │   │   │       │   │   │   │       │           kernel=SdpaFused.Config()
        │   │   │   │       │   │   │   │       │       )
        │   │   │   │       │   │   │   │       │   ),
        │   │   │   │       │   │   │   │       │   dropout=Dropout.Config(
        │   │   │   │       │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │       channels_out=1_024,
        │   │   │   │       │   │   │   │       │       p=0.1,
        │   │   │   │       │   │   │   │       │       inplace=False
        │   │   │   │       │   │   │   │       │   )
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       ffn=Sequential.Config(
        │   │   │   │       │   │   │   │       │   channels_in=1_024,
        │   │   │   │       │   │   │   │       │   channels_out=1_024,
        │   │   │   │       │   │   │   │       │   elements=[
        │   │   │   │       │   │   │   │       │   │   DreamFFN.Config(
        │   │   │   │       │   │   │   │       │   │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │       ffn=SwiGLU.Config(
        │   │   │   │       │   │   │   │       │   │           channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │           channels_out=1_024,
        │   │   │   │       │   │   │   │       │   │           channels_hidden=4_096,
        │   │   │   │       │   │   │   │       │   │           expansion=4.0,
        │   │   │   │       │   │   │   │       │   │           round_to=1,
        │   │   │   │       │   │   │   │       │   │           gate=False,
        │   │   │   │       │   │   │   │       │   │           bias=False,
        │   │   │   │       │   │   │   │       │   │           split_gate_projection=False,
        │   │   │   │       │   │   │   │       │   │           norm=None,
        │   │   │   │       │   │   │   │       │   │           act=priml.model.swiglu.relu,
        │   │   │   │       │   │   │   │       │   │           depth_index=(),
        │   │   │   │       │   │   │   │       │   │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │       │   │   │   │       │   │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │   │           up_proj=Linear.Config(
        │   │   │   │       │   │   │   │       │   │               channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │               channels_out=4_096,
        │   │   │   │       │   │   │   │       │   │               bias=False,
        │   │   │   │       │   │   │   │       │   │               depth_index=(),
        │   │   │   │       │   │   │   │       │   │               device=None,
        │   │   │   │       │   │   │   │       │   │               dtype=None,
        │   │   │   │       │   │   │   │       │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │   │       │   │               init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │   │               linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │   │               shard=None
        │   │   │   │       │   │   │   │       │   │           ),
        │   │   │   │       │   │   │   │       │   │           middle=Dropout.Config(
        │   │   │   │       │   │   │   │       │   │               channels_in=4_096,
        │   │   │   │       │   │   │   │       │   │               channels_out=4_096,
        │   │   │   │       │   │   │   │       │   │               p=0.1,
        │   │   │   │       │   │   │   │       │   │               inplace=False
        │   │   │   │       │   │   │   │       │   │           ),
        │   │   │   │       │   │   │   │       │   │           down_proj=Linear.Config(
        │   │   │   │       │   │   │   │       │   │               channels_in=4_096,
        │   │   │   │       │   │   │   │       │   │               channels_out=1_024,
        │   │   │   │       │   │   │   │       │   │               bias=False,
        │   │   │   │       │   │   │   │       │   │               depth_index=(),
        │   │   │   │       │   │   │   │       │   │               device=None,
        │   │   │   │       │   │   │   │       │   │               dtype=None,
        │   │   │   │       │   │   │   │       │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │   │       │   │               init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │   │               linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │   │               shard=None
        │   │   │   │       │   │   │   │       │   │           ),
        │   │   │   │       │   │   │   │       │   │           shard=None
        │   │   │   │       │   │   │   │       │   │       ),
        │   │   │   │       │   │   │   │       │   │       init_final=WidthNormal(width=1_024)
        │   │   │   │       │   │   │   │       │   │   ),
        │   │   │   │       │   │   │   │       │   │   Dropout.Config(
        │   │   │   │       │   │   │   │       │   │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │       channels_out=1_024,
        │   │   │   │       │   │   │   │       │   │       p=0.1,
        │   │   │   │       │   │   │   │       │   │       inplace=False
        │   │   │   │       │   │   │   │       │   │   )
        │   │   │   │       │   │   │   │       │   ],
        │   │   │   │       │   │   │   │       │   repeat=1,
        │   │   │   │       │   │   │   │       │   depth_index=((5, 7),)
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       norm1=LayerNorm.Config(
        │   │   │   │       │   │   │   │           channels_in=1_024,
        │   │   │   │       │   │   │   │           channels_out=1_024,
        │   │   │   │       │   │   │   │           eps=1e-05,
        │   │   │   │       │   │   │   │           elementwise_affine=True,
        │   │   │   │       │   │   │   │           device=None,
        │   │   │   │       │   │   │   │           dtype=None
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       norm2=LayerNorm.Config(
        │   │   │   │       │   │   │   │           channels_in=1_024,
        │   │   │   │       │   │   │   │           channels_out=1_024,
        │   │   │   │       │   │   │   │           eps=1e-05,
        │   │   │   │       │   │   │   │           elementwise_affine=True,
        │   │   │   │       │   │   │   │           device=None,
        │   │   │   │       │   │   │   │           dtype=None
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       prenorm=True,
        │   │   │   │       │   │   │   │       checkpoint=False,
        │   │   │   │       │   │   │   │       depth_index=((5, 7),)
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   TransformerBlock.Config(
        │   │   │   │       │   │   │   │       channels_in=1_024,
        │   │   │   │       │   │   │   │       channels_out=1_024,
        │   │   │   │       │   │   │   │       attn=DreamAttention.Config(
        │   │   │   │       │   │   │   │       │   channels_in=1_024,
        │   │   │   │       │   │   │   │       │   channels_out=1_024,
        │   │   │   │       │   │   │   │       │   attention=Attention.Config(
        │   │   │   │       │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │       channels_out=1_024,
        │   │   │   │       │   │   │   │       │       num_heads=8,
        │   │   │   │       │   │   │   │       │       channels_head=128,
        │   │   │   │       │   │   │   │       │       num_heads_kv=8,
        │   │   │   │       │   │   │   │       │       bias=False,
        │   │   │   │       │   │   │   │       │       dropout=0.1,
        │   │   │   │       │   │   │   │       │       causal=False,
        │   │   │   │       │   │   │   │       │       rope=None,
        │   │   │   │       │   │   │   │       │       rotate=None,
        │   │   │   │       │   │   │   │       │       norm_qk=None,
        │   │   │   │       │   │   │   │       │       share_qk_norm=True,
        │   │   │   │       │   │   │   │       │       norm_out=None,
        │   │   │   │       │   │   │   │       │       split_qkv_projection=False,
        │   │   │   │       │   │   │   │       │       proj_qkv=Linear.Config(
        │   │   │   │       │   │   │   │       │           channels_in=1_024,
        │   │   │   │       │   │   │   │       │           channels_out=3_072,
        │   │   │   │       │   │   │   │       │           bias=False,
        │   │   │   │       │   │   │   │       │           depth_index=(),
        │   │   │   │       │   │   │   │       │           device=None,
        │   │   │   │       │   │   │   │       │           dtype=None,
        │   │   │   │       │   │   │   │       │           init_weight=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │           shard=None
        │   │   │   │       │   │   │   │       │       ),
        │   │   │   │       │   │   │   │       │       proj_out=Linear.Config(
        │   │   │   │       │   │   │   │       │           channels_in=1_024,
        │   │   │   │       │   │   │   │       │           channels_out=1_024,
        │   │   │   │       │   │   │   │       │           bias=False,
        │   │   │   │       │   │   │   │       │           depth_index=(),
        │   │   │   │       │   │   │   │       │           device=None,
        │   │   │   │       │   │   │   │       │           dtype=None,
        │   │   │   │       │   │   │   │       │           init_weight=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │           linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │           shard=None
        │   │   │   │       │   │   │   │       │       ),
        │   │   │   │       │   │   │   │       │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │   │       │       depth_index=(),
        │   │   │   │       │   │   │   │       │       attn_kernel=DreamKernel.Config(
        │   │   │   │       │   │   │   │       │           dtype=torch.float32,
        │   │   │   │       │   │   │   │       │           kernel=SdpaFused.Config()
        │   │   │   │       │   │   │   │       │       )
        │   │   │   │       │   │   │   │       │   ),
        │   │   │   │       │   │   │   │       │   dropout=Dropout.Config(
        │   │   │   │       │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │       channels_out=1_024,
        │   │   │   │       │   │   │   │       │       p=0.1,
        │   │   │   │       │   │   │   │       │       inplace=False
        │   │   │   │       │   │   │   │       │   )
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       ffn=Sequential.Config(
        │   │   │   │       │   │   │   │       │   channels_in=1_024,
        │   │   │   │       │   │   │   │       │   channels_out=1_024,
        │   │   │   │       │   │   │   │       │   elements=[
        │   │   │   │       │   │   │   │       │   │   DreamFFN.Config(
        │   │   │   │       │   │   │   │       │   │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │       ffn=SwiGLU.Config(
        │   │   │   │       │   │   │   │       │   │           channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │           channels_out=1_024,
        │   │   │   │       │   │   │   │       │   │           channels_hidden=4_096,
        │   │   │   │       │   │   │   │       │   │           expansion=4.0,
        │   │   │   │       │   │   │   │       │   │           round_to=1,
        │   │   │   │       │   │   │   │       │   │           gate=False,
        │   │   │   │       │   │   │   │       │   │           bias=False,
        │   │   │   │       │   │   │   │       │   │           split_gate_projection=False,
        │   │   │   │       │   │   │   │       │   │           norm=None,
        │   │   │   │       │   │   │   │       │   │           act=priml.model.swiglu.relu,
        │   │   │   │       │   │   │   │       │   │           depth_index=(),
        │   │   │   │       │   │   │   │       │   │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │       │   │   │   │       │   │           init_weight_out=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │   │           up_proj=Linear.Config(
        │   │   │   │       │   │   │   │       │   │               channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │               channels_out=4_096,
        │   │   │   │       │   │   │   │       │   │               bias=False,
        │   │   │   │       │   │   │   │       │   │               depth_index=(),
        │   │   │   │       │   │   │   │       │   │               device=None,
        │   │   │   │       │   │   │   │       │   │               dtype=None,
        │   │   │   │       │   │   │   │       │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │   │       │   │               init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │   │               linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │   │               shard=None
        │   │   │   │       │   │   │   │       │   │           ),
        │   │   │   │       │   │   │   │       │   │           middle=Dropout.Config(
        │   │   │   │       │   │   │   │       │   │               channels_in=4_096,
        │   │   │   │       │   │   │   │       │   │               channels_out=4_096,
        │   │   │   │       │   │   │   │       │   │               p=0.1,
        │   │   │   │       │   │   │   │       │   │               inplace=False
        │   │   │   │       │   │   │   │       │   │           ),
        │   │   │   │       │   │   │   │       │   │           down_proj=Linear.Config(
        │   │   │   │       │   │   │   │       │   │               channels_in=4_096,
        │   │   │   │       │   │   │   │       │   │               channels_out=1_024,
        │   │   │   │       │   │   │   │       │   │               bias=False,
        │   │   │   │       │   │   │   │       │   │               depth_index=(),
        │   │   │   │       │   │   │   │       │   │               device=None,
        │   │   │   │       │   │   │   │       │   │               dtype=None,
        │   │   │   │       │   │   │   │       │   │               init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │   │       │   │               init_bias=torch.nn.init.zeros_,
        │   │   │   │       │   │   │   │       │   │               linear_fn=<built-in function linear>,
        │   │   │   │       │   │   │   │       │   │               shard=None
        │   │   │   │       │   │   │   │       │   │           ),
        │   │   │   │       │   │   │   │       │   │           shard=None
        │   │   │   │       │   │   │   │       │   │       ),
        │   │   │   │       │   │   │   │       │   │       init_final=WidthNormal(width=1_024)
        │   │   │   │       │   │   │   │       │   │   ),
        │   │   │   │       │   │   │   │       │   │   Dropout.Config(
        │   │   │   │       │   │   │   │       │   │       channels_in=1_024,
        │   │   │   │       │   │   │   │       │   │       channels_out=1_024,
        │   │   │   │       │   │   │   │       │   │       p=0.1,
        │   │   │   │       │   │   │   │       │   │       inplace=False
        │   │   │   │       │   │   │   │       │   │   )
        │   │   │   │       │   │   │   │       │   ],
        │   │   │   │       │   │   │   │       │   repeat=1,
        │   │   │   │       │   │   │   │       │   depth_index=((6, 7),)
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       norm1=LayerNorm.Config(
        │   │   │   │       │   │   │   │           channels_in=1_024,
        │   │   │   │       │   │   │   │           channels_out=1_024,
        │   │   │   │       │   │   │   │           eps=1e-05,
        │   │   │   │       │   │   │   │           elementwise_affine=True,
        │   │   │   │       │   │   │   │           device=None,
        │   │   │   │       │   │   │   │           dtype=None
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       norm2=LayerNorm.Config(
        │   │   │   │       │   │   │   │           channels_in=1_024,
        │   │   │   │       │   │   │   │           channels_out=1_024,
        │   │   │   │       │   │   │   │           eps=1e-05,
        │   │   │   │       │   │   │   │           elementwise_affine=True,
        │   │   │   │       │   │   │   │           device=None,
        │   │   │   │       │   │   │   │           dtype=None
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       prenorm=True,
        │   │   │   │       │   │   │   │       checkpoint=False,
        │   │   │   │       │   │   │   │       depth_index=((6, 7),)
        │   │   │   │       │   │   │   │   )
        │   │   │   │       │   │   │   ],
        │   │   │   │       │   │   │   proj_out=LayerNorm.Config(
        │   │   │   │       │   │   │       channels_in=1_024,
        │   │   │   │       │   │   │       channels_out=1_024,
        │   │   │   │       │   │   │       eps=1e-05,
        │   │   │   │       │   │   │       elementwise_affine=True,
        │   │   │   │       │   │   │       device=None,
        │   │   │   │       │   │   │       dtype=None
        │   │   │   │       │   │   │   )
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   chunk_size=4_096,
        │   │   │   │       │   │   checkpoint_fourier=True
        │   │   │   │       │   ),
        │   │   │   │       │   adduct=Embedding.Config(
        │   │   │   │       │       channels_in=60,
        │   │   │   │       │       channels_out=1_024,
        │   │   │   │       │       padding_idx=None,
        │   │   │   │       │       device=None,
        │   │   │   │       │       dtype=None,
        │   │   │   │       │       lookup_dtype=None,
        │   │   │   │       │       shard=None,
        │   │   │   │       │       depth_index=(),
        │   │   │   │       │       init_weight=torch.nn.init.normal_
        │   │   │   │       │   ),
        │   │   │   │       │   polarity=Embedding.Config(
        │   │   │   │       │       channels_in=2,
        │   │   │   │       │       channels_out=1_024,
        │   │   │   │       │       padding_idx=None,
        │   │   │   │       │       device=None,
        │   │   │   │       │       dtype=None,
        │   │   │   │       │       lookup_dtype=None,
        │   │   │   │       │       shard=None,
        │   │   │   │       │       depth_index=(),
        │   │   │   │       │       init_weight=torch.nn.init.normal_
        │   │   │   │       │   ),
        │   │   │   │       │   ce=SinEmb.Config(
        │   │   │   │       │       channels_out=32,
        │   │   │   │       │       lo=0.0,
        │   │   │   │       │       hi=2.2,
        │   │   │   │       │       power=1.0
        │   │   │   │       │   ),
        │   │   │   │       │   condition=Sequential.Config(
        │   │   │   │       │       channels_in=34,
        │   │   │   │       │       channels_out=1_024,
        │   │   │   │       │       elements=[
        │   │   │   │       │           Linear.Config(
        │   │   │   │       │               channels_in=34,
        │   │   │   │       │               channels_out=256,
        │   │   │   │       │               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
        │   │   │   │       │           ),
        │   │   │   │       │           Activation.Config(
        │   │   │   │       │               channels_in=256,
        │   │   │   │       │               channels_out=256,
        │   │   │   │       │               act=priml.model.swiglu.gelu
        │   │   │   │       │           ),
        │   │   │   │       │           Linear.Config(
        │   │   │   │       │               channels_in=256,
        │   │   │   │       │               channels_out=1_024,
        │   │   │   │       │               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=()
        │   │   │   │       │   ),
        │   │   │   │       │   head=Sequential.Config(
        │   │   │   │       │       channels_in=2_048,
        │   │   │   │       │       channels_out=10_226,
        │   │   │   │       │       elements=[
        │   │   │   │       │           Linear.Config(
        │   │   │   │       │               channels_in=2_048,
        │   │   │   │       │               channels_out=2_048,
        │   │   │   │       │               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
        │   │   │   │       │           ),
        │   │   │   │       │           Activation.Config(
        │   │   │   │       │               channels_in=2_048,
        │   │   │   │       │               channels_out=2_048,
        │   │   │   │       │               act=priml.model.swiglu.gelu
        │   │   │   │       │           ),
        │   │   │   │       │           Dropout.Config(
        │   │   │   │       │               channels_in=2_048,
        │   │   │   │       │               channels_out=2_048,
        │   │   │   │       │               p=0.1,
        │   │   │   │       │               inplace=False
        │   │   │   │       │           ),
        │   │   │   │       │           Linear.Config(
        │   │   │   │       │               channels_in=2_048,
        │   │   │   │       │               channels_out=10_226,
        │   │   │   │       │               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=()
        │   │   │   │       │   ),
        │   │   │   │       │   elements=Sequential.Config(
        │   │   │   │       │       channels_in=2_048,
        │   │   │   │       │       channels_out=10,
        │   │   │   │       │       elements=[
        │   │   │   │       │           Linear.Config(
        │   │   │   │       │               channels_in=2_048,
        │   │   │   │       │               channels_out=256,
        │   │   │   │       │               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
        │   │   │   │       │           ),
        │   │   │   │       │           Activation.Config(
        │   │   │   │       │               channels_in=256,
        │   │   │   │       │               channels_out=256,
        │   │   │   │       │               act=priml.model.swiglu.gelu
        │   │   │   │       │           ),
        │   │   │   │       │           Linear.Config(
        │   │   │   │       │               channels_in=256,
        │   │   │   │       │               channels_out=10,
        │   │   │   │       │               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=()
        │   │   │   │       │   )
        │   │   │   │       ),
        │   │   │   │       dtype_autocast=torch.bfloat16
        │   │   │   │   )
        │   │   │   ],
        │   │   │   device='cuda',
        │   │   │   batch_size=64,
        │   │   │   max_peaks=160,
        │   │   │   floor=0.001,
        │   │   │   merge_tolerance=0.005,
        │   │   │   merged_adduct=priml.baselines.casmi2026.dreams.merged_adduct,
        │   │   │   single_count=priml.baselines.casmi2026.dreams.clamp_count
        │   │   )
        │   ),
        │   engine=MainEngine.Config(
        │       retrieval=MainRetrieval.Config(
        │           search=MainSearch.Config(
        │               floor=0.002,
        │               max_peaks=256,
        │               power=1.0,
        │               entropy_weight=True,
        │               tolerance_da=0.01,
        │               tolerance_ppm=20.0
        │           ),
        │           analog_window=200.0,
        │           n_analog=100,
        │           merge=PolarityMerge.Config(tolerance=0.005)
        │       ),
        │       max_query_spectra=16,
        │       ppm=10.0,
        │       fallback_ppm=30.0,
        │       similarity_power=4.0,
        │       fragment_search=MainSearch.Config(
        │           floor=0.002,
        │           max_peaks=128,
        │           power=1.0,
        │           entropy_weight=False,
        │           tolerance_da=0.01,
        │           tolerance_ppm=20.0
        │       ),
        │       fragments=FragmentChannels.Config(max_candidates=600),
        │       subformulas=SubformulaChannels.Config(max_out=2000000),
        │       query_floor=0.001,
        │       query_max_peaks=512,
        │       query_over=2.0,
        │       family_max_candidates=600,
        │       pairwise_rank_limit=3_000
        │   ),
        │   generation=ReactionScreen.Config(
        │       parents=6,
        │       min_similarity=0.3,
        │       per_parent=60,
        │       total=150,
        │       tolerance=0.004,
        │       ppm=10.0,
        │       max_steps=2,
        │       application_budget=800,
        │       stage_parents=24,
        │       first_products=32,
        │       later_products=16
        │   ),
        │   limit=60,
        │   popularity_weight=0.15
        ),
        second=SecondResource.Config(
        │   assets=SecondAssets.Config(
        │   │   base_dir=PosixPath('/opt/scratch/datasets/casmi2026'),
        │   │   pool_dir=PosixPath('/opt/scratch/datasets/casmi2026/native-chemistry/engine2-pool'),
        │   │   library_dir=PosixPath('/opt/scratch/datasets/casmi2026/native-chemistry/engine2-library'),
        │   │   fragments_dir=PosixPath('/opt/scratch/datasets/casmi2026/native-chemistry/engine2-fragments'),
        │   │   chemistry=IdentityTable.Config(
        │   │       base_dir=PosixPath('/opt/scratch/datasets/casmi2026'),
        │   │       working_dir=PosixPath('/opt/scratch/datasets/casmi2026/native-chemistry/identities')
        │   │   ),
        │   │   identities_dir=PosixPath('/opt/scratch/datasets/casmi2026/native-chemistry/engine-identities'),
        │   │   own_training=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/megayak/casmi26-simulated-ranker-rows/sim_rank_rows_nofp.npz'),
        │   │   pv_training=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/prvsiyan/casmi26-ranker-features/rank_train.npz'),
        │   │   bank=SecondBank.Config(
        │   │   │   base_dir=PosixPath('/opt/scratch/datasets/casmi2026'),
        │   │   │   model=SecondFPNet.Config(
        │   │   │   │   channels_hidden=512,
        │   │   │   │   channels_out=6_930,
        │   │   │   │   num_layers=6,
        │   │   │   │   mz_emb=SinEmb.Config(
        │   │   │   │       channels_out=512,
        │   │   │   │       lo=-2.0,
        │   │   │   │       hi=3.2,
        │   │   │   │       power=1.0
        │   │   │   │   ),
        │   │   │   │   nl_emb=SinEmb.Config(
        │   │   │   │       channels_out=512,
        │   │   │   │       lo=-2.0,
        │   │   │   │       hi=3.2,
        │   │   │   │       power=1.0
        │   │   │   │   ),
        │   │   │   │   prec_emb=SinEmb.Config(
        │   │   │   │       channels_out=512,
        │   │   │   │       lo=-2.0,
        │   │   │   │       hi=3.2,
        │   │   │   │       power=1.0
        │   │   │   │   ),
        │   │   │   │   peak_projection=Linear.Config(
        │   │   │   │       channels_in=1_025,
        │   │   │   │       channels_out=512,
        │   │   │   │       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
        │   │   │   │   ),
        │   │   │   │   adduct=Embedding.Config(
        │   │   │   │       channels_in=26,
        │   │   │   │       channels_out=512,
        │   │   │   │       padding_idx=None,
        │   │   │   │       device=None,
        │   │   │   │       dtype=None,
        │   │   │   │       lookup_dtype=None,
        │   │   │   │       shard=None,
        │   │   │   │       depth_index=(),
        │   │   │   │       init_weight=torch.nn.init.normal_
        │   │   │   │   ),
        │   │   │   │   instrument=Embedding.Config(
        │   │   │   │       channels_in=5,
        │   │   │   │       channels_out=512,
        │   │   │   │       padding_idx=None,
        │   │   │   │       device=None,
        │   │   │   │       dtype=None,
        │   │   │   │       lookup_dtype=None,
        │   │   │   │       shard=None,
        │   │   │   │       depth_index=(),
        │   │   │   │       init_weight=torch.nn.init.normal_
        │   │   │   │   ),
        │   │   │   │   context_projection=Linear.Config(
        │   │   │   │       channels_in=515,
        │   │   │   │       channels_out=512,
        │   │   │   │       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
        │   │   │   │   ),
        │   │   │   │   block=Block.Config(
        │   │   │   │   │   channels_in=512,
        │   │   │   │   │   channels_out=512,
        │   │   │   │   │   attn=PeakAttention.Config(
        │   │   │   │   │   │   channels_in=512,
        │   │   │   │   │   │   channels_out=512,
        │   │   │   │   │   │   attention=Attention.Config(
        │   │   │   │   │   │       channels_in=512,
        │   │   │   │   │   │       channels_out=512,
        │   │   │   │   │   │       num_heads=8,
        │   │   │   │   │   │       channels_head=64,
        │   │   │   │   │   │       num_heads_kv=8,
        │   │   │   │   │   │       bias=False,
        │   │   │   │   │   │       dropout=0.0,
        │   │   │   │   │   │       causal=False,
        │   │   │   │   │   │       rope=None,
        │   │   │   │   │   │       rotate=None,
        │   │   │   │   │   │       norm_qk=None,
        │   │   │   │   │   │       share_qk_norm=True,
        │   │   │   │   │   │       norm_out=None,
        │   │   │   │   │   │       split_qkv_projection=False,
        │   │   │   │   │   │       proj_qkv=Linear.Config(
        │   │   │   │   │   │           channels_in=512,
        │   │   │   │   │   │           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
        │   │   │   │   │   │       ),
        │   │   │   │   │   │       proj_out=Linear.Config(
        │   │   │   │   │   │           channels_in=512,
        │   │   │   │   │   │           channels_out=512,
        │   │   │   │   │   │           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
        │   │   │   │   │   │       ),
        │   │   │   │   │   │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │       depth_index=(),
        │   │   │   │   │   │       attn_kernel=SdpaFused.Config()
        │   │   │   │   │   │   ),
        │   │   │   │   │   │   drop=Dropout.Config(
        │   │   │   │   │   │       channels_in=512,
        │   │   │   │   │   │       channels_out=512,
        │   │   │   │   │   │       p=0.1,
        │   │   │   │   │   │       inplace=False
        │   │   │   │   │   │   )
        │   │   │   │   │   ),
        │   │   │   │   │   ffn=Sequential.Config(
        │   │   │   │   │   │   channels_in=512,
        │   │   │   │   │   │   channels_out=512,
        │   │   │   │   │   │   elements=[
        │   │   │   │   │   │   │   SwiGLU.Config(
        │   │   │   │   │   │   │       channels_in=512,
        │   │   │   │   │   │   │       channels_out=512,
        │   │   │   │   │   │   │       channels_hidden=2_048,
        │   │   │   │   │   │   │       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=512,
        │   │   │   │   │   │   │           channels_out=2_048,
        │   │   │   │   │   │   │           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=Dropout.Config(
        │   │   │   │   │   │   │           channels_in=2_048,
        │   │   │   │   │   │   │           channels_out=2_048,
        │   │   │   │   │   │   │           p=0.1,
        │   │   │   │   │   │   │           inplace=False
        │   │   │   │   │   │   │       ),
        │   │   │   │   │   │   │       down_proj=Linear.Config(
        │   │   │   │   │   │   │           channels_in=2_048,
        │   │   │   │   │   │   │           channels_out=512,
        │   │   │   │   │   │   │           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
        │   │   │   │   │   │   │   ),
        │   │   │   │   │   │   │   Dropout.Config(
        │   │   │   │   │   │   │       channels_in=512,
        │   │   │   │   │   │   │       channels_out=512,
        │   │   │   │   │   │   │       p=0.1,
        │   │   │   │   │   │   │       inplace=False
        │   │   │   │   │   │   │   )
        │   │   │   │   │   │   ],
        │   │   │   │   │   │   repeat=1,
        │   │   │   │   │   │   depth_index=()
        │   │   │   │   │   ),
        │   │   │   │   │   norm1=LayerNorm.Config(
        │   │   │   │   │       channels_in=512,
        │   │   │   │   │       channels_out=512,
        │   │   │   │   │       eps=1e-05,
        │   │   │   │   │       elementwise_affine=True,
        │   │   │   │   │       device=None,
        │   │   │   │   │       dtype=None
        │   │   │   │   │   ),
        │   │   │   │   │   norm2=LayerNorm.Config(
        │   │   │   │   │       channels_in=512,
        │   │   │   │   │       channels_out=512,
        │   │   │   │   │       eps=1e-05,
        │   │   │   │   │       elementwise_affine=True,
        │   │   │   │   │       device=None,
        │   │   │   │   │       dtype=None
        │   │   │   │   │   ),
        │   │   │   │   │   prenorm=True,
        │   │   │   │   │   checkpoint=False,
        │   │   │   │   │   depth_index=()
        │   │   │   │   ),
        │   │   │   │   norm=LayerNorm.Config(
        │   │   │   │       channels_in=512,
        │   │   │   │       channels_out=512,
        │   │   │   │       eps=1e-05,
        │   │   │   │       elementwise_affine=True,
        │   │   │   │       device=None,
        │   │   │   │       dtype=None
        │   │   │   │   ),
        │   │   │   │   head=SwiGLU.Config(
        │   │   │   │       channels_in=1_024,
        │   │   │   │       channels_out=6_930,
        │   │   │   │       channels_hidden=2_048,
        │   │   │   │       expansion=4.0,
        │   │   │   │       round_to=256,
        │   │   │   │       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=1_024,
        │   │   │   │           channels_out=2_048,
        │   │   │   │           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=Dropout.Config(
        │   │   │   │           channels_in=2_048,
        │   │   │   │           channels_out=2_048,
        │   │   │   │           p=0.1,
        │   │   │   │           inplace=False
        │   │   │   │       ),
        │   │   │   │       down_proj=Linear.Config(
        │   │   │   │           channels_in=2_048,
        │   │   │   │           channels_out=6_930,
        │   │   │   │           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
        │   │   │   │   )
        │   │   │   ),
        │   │   │   single=[
        │   │   │       LocalFile.Config(
        │   │   │           path=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/prvsiyan/casmi26-fp-models-v2/fp_single_s2.pt'),
        │   │   │           sha256='c19aa29016ed4d7a67bac38105b8fa5076c696c29eb66a504558c9b2e52b614a'
        │   │   │       )
        │   │   │   ],
        │   │   │   merged=[
        │   │   │       LocalFile.Config(
        │   │   │           path=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/prvsiyan/casmi26-fp-models-v2/fp_merged_m1.pt'),
        │   │   │           sha256='261529c823825cbe2028684bbf7214a95a7e631ca9f656336d577c507cce8294'
        │   │   │       )
        │   │   │   ],
        │   │   │   peaks=SecondPeaks.Config(
        │   │   │       max_peaks=128,
        │   │   │       floor=0.001,
        │   │   │       window=50.0,
        │   │   │       per_window=8,
        │   │   │       margin=1.5,
        │   │   │       merge_tolerance=0.005
        │   │   │   ),
        │   │   │   merged_energy=25.0,
        │   │   │   device='cuda'
        │   │   ),
        │   │   rankers=SecondRankerBank.Config(
        │   │       own=HistogramEnsemble.Config(
        │   │           priors=[0.35, 0.55],
        │   │           seeds=[0, 1],
        │   │           tree=HistogramBoosting.Config(
        │   │               max_depth=6,
        │   │               rounds=500,
        │   │               learning_rate=0.03,
        │   │               min_leaf=80,
        │   │               l2=1.0,
        │   │               max_bins=255,
        │   │               max_leaves=31,
        │   │               seed=-1,
        │   │               early_stopping_rows=10_000,
        │   │               binning_rows=200_000,
        │   │               validation_fraction=0.1,
        │   │               patience=10,
        │   │               tolerance=1e-07
        │   │           )
        │   │       ),
        │   │       pv=HistogramEnsemble.Config(
        │   │           priors=[0.3, 0.6],
        │   │           seeds=[0, 1, 2, 3],
        │   │           tree=HistogramBoosting.Config(
        │   │               max_depth=6,
        │   │               rounds=500,
        │   │               learning_rate=0.03,
        │   │               min_leaf=80,
        │   │               l2=1.0,
        │   │               max_bins=255,
        │   │               max_leaves=31,
        │   │               seed=-1,
        │   │               early_stopping_rows=10_000,
        │   │               binning_rows=200_000,
        │   │               validation_fraction=0.1,
        │   │               patience=10,
        │   │               tolerance=1e-07
        │   │           )
        │   │       ),
        │   │       weight_fingerprint=0.88
        │   │   )
        │   ),
        │   engine=SecondEngine.Config(
        │       retrieval=SecondRetrieval.Config(
        │           search=SecondSearch.Config(
        │               floor=0.002,
        │               max_peaks=256,
        │               power=1.0,
        │               entropy_weight=True,
        │               tolerance_da=0.01
        │           ),
        │           analog_window=200.0,
        │           ppm=10.0,
        │           n_analog=200
        │       ),
        │       ppm=10.0,
        │       fallback_ppm=30.0,
        │       precursor_margin=2.0,
        │       default_energy=25.0,
        │       fragment_search=SecondSearch.Config(
        │           floor=0.002,
        │           max_peaks=256,
        │           power=1.0,
        │           entropy_weight=False,
        │           tolerance_da=0.01
        │       )
        │   ),
        │   prescreen=80,
        │   limit=40
        ),
        pubchem=PubChemResource.Config(
        │   assets=PubChemAssets.Config(
        │   │   base_dir=PosixPath('/opt/scratch/datasets/casmi2026'),
        │   │   pool_dir=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/ahmedberatozer/casmi26-v2-pool'),
        │   │   store_dir=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/ahmedberatozer/casmi26-pubchem-tier'),
        │   │   popularity_dir=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/dmitriigluzdov/casmi26-pubchem-popularity-prior'),
        │   │   chemistry_dir=PosixPath('/opt/scratch/datasets/casmi2026/native-chemistry/pubchem'),
        │   │   bank=MainBank.Config(
        │   │   │   checkpoints=[
        │   │   │   │   BankCheckpoint.Config(
        │   │   │   │       base_dir=PosixPath('/opt/scratch/datasets/casmi2026'),
        │   │   │   │       path=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/ahmedberatozer/casmi26-v3-models/fpnet_0.pt'),
        │   │   │   │       sha256='edccf479a2a2e53461a06bc5e1aee4341b22abb32302c4f6d2c45281000d742a',
        │   │   │   │       model=FPNet.Config(
        │   │   │   │       │   channels_out=10_226,
        │   │   │   │       │   channels_hidden=512,
        │   │   │   │       │   num_layers=8,
        │   │   │   │       │   block=Block.Config(
        │   │   │   │       │   │   channels_in=512,
        │   │   │   │       │   │   channels_out=512,
        │   │   │   │       │   │   attn=PeakAttention.Config(
        │   │   │   │       │   │   │   channels_in=512,
        │   │   │   │       │   │   │   channels_out=512,
        │   │   │   │       │   │   │   attention=Attention.Config(
        │   │   │   │       │   │   │       channels_in=512,
        │   │   │   │       │   │   │       channels_out=512,
        │   │   │   │       │   │   │       num_heads=8,
        │   │   │   │       │   │   │       channels_head=64,
        │   │   │   │       │   │   │       num_heads_kv=8,
        │   │   │   │       │   │   │       bias=False,
        │   │   │   │       │   │   │       dropout=0.0,
        │   │   │   │       │   │   │       causal=False,
        │   │   │   │       │   │   │       rope=None,
        │   │   │   │       │   │   │       rotate=None,
        │   │   │   │       │   │   │       norm_qk=None,
        │   │   │   │       │   │   │       share_qk_norm=True,
        │   │   │   │       │   │   │       norm_out=None,
        │   │   │   │       │   │   │       split_qkv_projection=False,
        │   │   │   │       │   │   │       proj_qkv=Linear.Config(
        │   │   │   │       │   │   │           channels_in=512,
        │   │   │   │       │   │   │           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
        │   │   │   │       │   │   │       ),
        │   │   │   │       │   │   │       proj_out=Linear.Config(
        │   │   │   │       │   │   │           channels_in=512,
        │   │   │   │       │   │   │           channels_out=512,
        │   │   │   │       │   │   │           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
        │   │   │   │       │   │   │       ),
        │   │   │   │       │   │   │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │       depth_index=(),
        │   │   │   │       │   │   │       attn_kernel=SdpaFused.Config()
        │   │   │   │       │   │   │   ),
        │   │   │   │       │   │   │   drop=Dropout.Config(
        │   │   │   │       │   │   │       channels_in=512,
        │   │   │   │       │   │   │       channels_out=512,
        │   │   │   │       │   │   │       p=0.1,
        │   │   │   │       │   │   │       inplace=False
        │   │   │   │       │   │   │   )
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   ffn=Sequential.Config(
        │   │   │   │       │   │   │   channels_in=512,
        │   │   │   │       │   │   │   channels_out=512,
        │   │   │   │       │   │   │   elements=[
        │   │   │   │       │   │   │   │   SwiGLU.Config(
        │   │   │   │       │   │   │   │       channels_in=512,
        │   │   │   │       │   │   │   │       channels_out=512,
        │   │   │   │       │   │   │   │       channels_hidden=2_048,
        │   │   │   │       │   │   │   │       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=512,
        │   │   │   │       │   │   │   │           channels_out=2_048,
        │   │   │   │       │   │   │   │           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=Dropout.Config(
        │   │   │   │       │   │   │   │           channels_in=2_048,
        │   │   │   │       │   │   │   │           channels_out=2_048,
        │   │   │   │       │   │   │   │           p=0.1,
        │   │   │   │       │   │   │   │           inplace=False
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       down_proj=Linear.Config(
        │   │   │   │       │   │   │   │           channels_in=2_048,
        │   │   │   │       │   │   │   │           channels_out=512,
        │   │   │   │       │   │   │   │           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
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   Dropout.Config(
        │   │   │   │       │   │   │   │       channels_in=512,
        │   │   │   │       │   │   │   │       channels_out=512,
        │   │   │   │       │   │   │   │       p=0.1,
        │   │   │   │       │   │   │   │       inplace=False
        │   │   │   │       │   │   │   │   )
        │   │   │   │       │   │   │   ],
        │   │   │   │       │   │   │   repeat=1,
        │   │   │   │       │   │   │   depth_index=()
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   norm1=LayerNorm.Config(
        │   │   │   │       │   │       channels_in=512,
        │   │   │   │       │   │       channels_out=512,
        │   │   │   │       │   │       eps=1e-05,
        │   │   │   │       │   │       elementwise_affine=True,
        │   │   │   │       │   │       device=None,
        │   │   │   │       │   │       dtype=None
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   norm2=LayerNorm.Config(
        │   │   │   │       │   │       channels_in=512,
        │   │   │   │       │   │       channels_out=512,
        │   │   │   │       │   │       eps=1e-05,
        │   │   │   │       │   │       elementwise_affine=True,
        │   │   │   │       │   │       device=None,
        │   │   │   │       │   │       dtype=None
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   prenorm=True,
        │   │   │   │       │   │   checkpoint=False,
        │   │   │   │       │   │   depth_index=()
        │   │   │   │       │   ),
        │   │   │   │       │   mz_emb=SinEmb.Config(
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       lo=-2.0,
        │   │   │   │       │       hi=3.3,
        │   │   │   │       │       power=1.0
        │   │   │   │       │   ),
        │   │   │   │       │   nl_emb=SinEmb.Config(
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       lo=-2.0,
        │   │   │   │       │       hi=3.3,
        │   │   │   │       │       power=1.0
        │   │   │   │       │   ),
        │   │   │   │       │   prec_emb=SinEmb.Config(
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       lo=-2.0,
        │   │   │   │       │       hi=3.3,
        │   │   │   │       │       power=1.0
        │   │   │   │       │   ),
        │   │   │   │       │   ce_emb=SinEmb.Config(
        │   │   │   │       │       channels_out=128,
        │   │   │   │       │       lo=0.0,
        │   │   │   │       │       hi=2.2,
        │   │   │   │       │       power=1.0
        │   │   │   │       │   ),
        │   │   │   │       │   int_emb=Linear.Config(
        │   │   │   │       │       channels_in=2,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       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
        │   │   │   │       │   ),
        │   │   │   │       │   pk=Linear.Config(
        │   │   │   │       │       channels_in=1_536,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       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
        │   │   │   │       │   ),
        │   │   │   │       │   adduct=Embedding.Config(
        │   │   │   │       │       channels_in=60,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       padding_idx=None,
        │   │   │   │       │       device=None,
        │   │   │   │       │       dtype=None,
        │   │   │   │       │       lookup_dtype=None,
        │   │   │   │       │       shard=None,
        │   │   │   │       │       depth_index=(),
        │   │   │   │       │       init_weight=torch.nn.init.normal_
        │   │   │   │       │   ),
        │   │   │   │       │   gl=Linear.Config(
        │   │   │   │       │       channels_in=644,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       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
        │   │   │   │       │   ),
        │   │   │   │       │   norm=LayerNorm.Config(
        │   │   │   │       │       channels_in=512,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       eps=1e-05,
        │   │   │   │       │       elementwise_affine=True,
        │   │   │   │       │       device=None,
        │   │   │   │       │       dtype=None
        │   │   │   │       │   ),
        │   │   │   │       │   head=SwiGLU.Config(
        │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │       channels_out=10_226,
        │   │   │   │       │       channels_hidden=2_048,
        │   │   │   │       │       expansion=4.0,
        │   │   │   │       │       round_to=256,
        │   │   │   │       │       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=1_024,
        │   │   │   │       │           channels_out=2_048,
        │   │   │   │       │           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=Dropout.Config(
        │   │   │   │       │           channels_in=2_048,
        │   │   │   │       │           channels_out=2_048,
        │   │   │   │       │           p=0.1,
        │   │   │   │       │           inplace=False
        │   │   │   │       │       ),
        │   │   │   │       │       down_proj=Linear.Config(
        │   │   │   │       │           channels_in=2_048,
        │   │   │   │       │           channels_out=10_226,
        │   │   │   │       │           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
        │   │   │   │       │   ),
        │   │   │   │       │   head_elements=SwiGLU.Config(
        │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │       channels_out=10,
        │   │   │   │       │       channels_hidden=256,
        │   │   │   │       │       expansion=4.0,
        │   │   │   │       │       round_to=256,
        │   │   │   │       │       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=1_024,
        │   │   │   │       │           channels_out=256,
        │   │   │   │       │           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=256,
        │   │   │   │       │           channels_out=10,
        │   │   │   │       │           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
        │   │   │   │       │   )
        │   │   │   │       ),
        │   │   │   │       dtype_autocast=torch.bfloat16
        │   │   │   │   ),
        │   │   │   │   BankCheckpoint.Config(
        │   │   │   │       base_dir=PosixPath('/opt/scratch/datasets/casmi2026'),
        │   │   │   │       path=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/ahmedberatozer/casmi26-v3-models/fpnet_1.pt'),
        │   │   │   │       sha256='6f60a22b3c3b21aa48d47328bb9523bdf07cc004bfb72966e0f9713507c24f2e',
        │   │   │   │       model=FPNet.Config(
        │   │   │   │       │   channels_out=10_226,
        │   │   │   │       │   channels_hidden=512,
        │   │   │   │       │   num_layers=8,
        │   │   │   │       │   block=Block.Config(
        │   │   │   │       │   │   channels_in=512,
        │   │   │   │       │   │   channels_out=512,
        │   │   │   │       │   │   attn=PeakAttention.Config(
        │   │   │   │       │   │   │   channels_in=512,
        │   │   │   │       │   │   │   channels_out=512,
        │   │   │   │       │   │   │   attention=Attention.Config(
        │   │   │   │       │   │   │       channels_in=512,
        │   │   │   │       │   │   │       channels_out=512,
        │   │   │   │       │   │   │       num_heads=8,
        │   │   │   │       │   │   │       channels_head=64,
        │   │   │   │       │   │   │       num_heads_kv=8,
        │   │   │   │       │   │   │       bias=False,
        │   │   │   │       │   │   │       dropout=0.0,
        │   │   │   │       │   │   │       causal=False,
        │   │   │   │       │   │   │       rope=None,
        │   │   │   │       │   │   │       rotate=None,
        │   │   │   │       │   │   │       norm_qk=None,
        │   │   │   │       │   │   │       share_qk_norm=True,
        │   │   │   │       │   │   │       norm_out=None,
        │   │   │   │       │   │   │       split_qkv_projection=False,
        │   │   │   │       │   │   │       proj_qkv=Linear.Config(
        │   │   │   │       │   │   │           channels_in=512,
        │   │   │   │       │   │   │           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
        │   │   │   │       │   │   │       ),
        │   │   │   │       │   │   │       proj_out=Linear.Config(
        │   │   │   │       │   │   │           channels_in=512,
        │   │   │   │       │   │   │           channels_out=512,
        │   │   │   │       │   │   │           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
        │   │   │   │       │   │   │       ),
        │   │   │   │       │   │   │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │       │   │   │       depth_index=(),
        │   │   │   │       │   │   │       attn_kernel=SdpaFused.Config()
        │   │   │   │       │   │   │   ),
        │   │   │   │       │   │   │   drop=Dropout.Config(
        │   │   │   │       │   │   │       channels_in=512,
        │   │   │   │       │   │   │       channels_out=512,
        │   │   │   │       │   │   │       p=0.1,
        │   │   │   │       │   │   │       inplace=False
        │   │   │   │       │   │   │   )
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   ffn=Sequential.Config(
        │   │   │   │       │   │   │   channels_in=512,
        │   │   │   │       │   │   │   channels_out=512,
        │   │   │   │       │   │   │   elements=[
        │   │   │   │       │   │   │   │   SwiGLU.Config(
        │   │   │   │       │   │   │   │       channels_in=512,
        │   │   │   │       │   │   │   │       channels_out=512,
        │   │   │   │       │   │   │   │       channels_hidden=2_048,
        │   │   │   │       │   │   │   │       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=512,
        │   │   │   │       │   │   │   │           channels_out=2_048,
        │   │   │   │       │   │   │   │           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=Dropout.Config(
        │   │   │   │       │   │   │   │           channels_in=2_048,
        │   │   │   │       │   │   │   │           channels_out=2_048,
        │   │   │   │       │   │   │   │           p=0.1,
        │   │   │   │       │   │   │   │           inplace=False
        │   │   │   │       │   │   │   │       ),
        │   │   │   │       │   │   │   │       down_proj=Linear.Config(
        │   │   │   │       │   │   │   │           channels_in=2_048,
        │   │   │   │       │   │   │   │           channels_out=512,
        │   │   │   │       │   │   │   │           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
        │   │   │   │       │   │   │   │   ),
        │   │   │   │       │   │   │   │   Dropout.Config(
        │   │   │   │       │   │   │   │       channels_in=512,
        │   │   │   │       │   │   │   │       channels_out=512,
        │   │   │   │       │   │   │   │       p=0.1,
        │   │   │   │       │   │   │   │       inplace=False
        │   │   │   │       │   │   │   │   )
        │   │   │   │       │   │   │   ],
        │   │   │   │       │   │   │   repeat=1,
        │   │   │   │       │   │   │   depth_index=()
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   norm1=LayerNorm.Config(
        │   │   │   │       │   │       channels_in=512,
        │   │   │   │       │   │       channels_out=512,
        │   │   │   │       │   │       eps=1e-05,
        │   │   │   │       │   │       elementwise_affine=True,
        │   │   │   │       │   │       device=None,
        │   │   │   │       │   │       dtype=None
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   norm2=LayerNorm.Config(
        │   │   │   │       │   │       channels_in=512,
        │   │   │   │       │   │       channels_out=512,
        │   │   │   │       │   │       eps=1e-05,
        │   │   │   │       │   │       elementwise_affine=True,
        │   │   │   │       │   │       device=None,
        │   │   │   │       │   │       dtype=None
        │   │   │   │       │   │   ),
        │   │   │   │       │   │   prenorm=True,
        │   │   │   │       │   │   checkpoint=False,
        │   │   │   │       │   │   depth_index=()
        │   │   │   │       │   ),
        │   │   │   │       │   mz_emb=SinEmb.Config(
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       lo=-2.0,
        │   │   │   │       │       hi=3.3,
        │   │   │   │       │       power=1.0
        │   │   │   │       │   ),
        │   │   │   │       │   nl_emb=SinEmb.Config(
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       lo=-2.0,
        │   │   │   │       │       hi=3.3,
        │   │   │   │       │       power=1.0
        │   │   │   │       │   ),
        │   │   │   │       │   prec_emb=SinEmb.Config(
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       lo=-2.0,
        │   │   │   │       │       hi=3.3,
        │   │   │   │       │       power=1.0
        │   │   │   │       │   ),
        │   │   │   │       │   ce_emb=SinEmb.Config(
        │   │   │   │       │       channels_out=128,
        │   │   │   │       │       lo=0.0,
        │   │   │   │       │       hi=2.2,
        │   │   │   │       │       power=1.0
        │   │   │   │       │   ),
        │   │   │   │       │   int_emb=Linear.Config(
        │   │   │   │       │       channels_in=2,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       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
        │   │   │   │       │   ),
        │   │   │   │       │   pk=Linear.Config(
        │   │   │   │       │       channels_in=1_536,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       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
        │   │   │   │       │   ),
        │   │   │   │       │   adduct=Embedding.Config(
        │   │   │   │       │       channels_in=60,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       padding_idx=None,
        │   │   │   │       │       device=None,
        │   │   │   │       │       dtype=None,
        │   │   │   │       │       lookup_dtype=None,
        │   │   │   │       │       shard=None,
        │   │   │   │       │       depth_index=(),
        │   │   │   │       │       init_weight=torch.nn.init.normal_
        │   │   │   │       │   ),
        │   │   │   │       │   gl=Linear.Config(
        │   │   │   │       │       channels_in=644,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       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
        │   │   │   │       │   ),
        │   │   │   │       │   norm=LayerNorm.Config(
        │   │   │   │       │       channels_in=512,
        │   │   │   │       │       channels_out=512,
        │   │   │   │       │       eps=1e-05,
        │   │   │   │       │       elementwise_affine=True,
        │   │   │   │       │       device=None,
        │   │   │   │       │       dtype=None
        │   │   │   │       │   ),
        │   │   │   │       │   head=SwiGLU.Config(
        │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │       channels_out=10_226,
        │   │   │   │       │       channels_hidden=2_048,
        │   │   │   │       │       expansion=4.0,
        │   │   │   │       │       round_to=256,
        │   │   │   │       │       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=1_024,
        │   │   │   │       │           channels_out=2_048,
        │   │   │   │       │           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=Dropout.Config(
        │   │   │   │       │           channels_in=2_048,
        │   │   │   │       │           channels_out=2_048,
        │   │   │   │       │           p=0.1,
        │   │   │   │       │           inplace=False
        │   │   │   │       │       ),
        │   │   │   │       │       down_proj=Linear.Config(
        │   │   │   │       │           channels_in=2_048,
        │   │   │   │       │           channels_out=10_226,
        │   │   │   │       │           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
        │   │   │   │       │   ),
        │   │   │   │       │   head_elements=SwiGLU.Config(
        │   │   │   │       │       channels_in=1_024,
        │   │   │   │       │       channels_out=10,
        │   │   │   │       │       channels_hidden=256,
        │   │   │   │       │       expansion=4.0,
        │   │   │   │       │       round_to=256,
        │   │   │   │       │       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=1_024,
        │   │   │   │       │           channels_out=256,
        │   │   │   │       │           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=256,
        │   │   │   │       │           channels_out=10,
        │   │   │   │       │           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
        │   │   │   │       │   )
        │   │   │   │       ),
        │   │   │   │       dtype_autocast=torch.bfloat16
        │   │   │   │   )
        │   │   │   ],
        │   │   │   device='cuda',
        │   │   │   batch_size=64,
        │   │   │   max_peaks=160,
        │   │   │   floor=0.001,
        │   │   │   merge_tolerance=0.005,
        │   │   │   merged_adduct=priml.baselines.casmi2026.pubchem.first_modal_adduct,
        │   │   │   single_count=priml.baselines.casmi2026.dreams.keep_count
        │   │   )
        │   ),
        │   screen=PubChemScreen.Config(
        │       pass_one=5_000,
        │       popularity_weight=0.25,
        │       popularity_union=200,
        │       limit=25,
        │       max_spectra=16,
        │       ppm=10.0,
        │       fallback_ppm=30.0
        │   )
        ),
        ice=SimulatorStage.Config(
        │   base_dir=PosixPath('/opt/scratch/datasets/casmi2026'),
        │   admission_dir=PosixPath('/opt/scratch/datasets/casmi2026/native-chemistry/simulator-admission'),
        │   backend=IceBackend.Config(
        │   │   base_dir=PosixPath('/opt/scratch/datasets/casmi2026'),
        │   │   root_dir=PosixPath('/opt/scratch/datasets/casmi2026/native-chemistry/ice-roots'),
        │   │   generator=GraphGenerator.Config(
        │   │   │   channels_in=46,
        │   │   │   channels_hidden=512,
        │   │   │   num_elements=18,
        │   │   │   max_broken=6,
        │   │   │   formula=FormulaEncoding.Config(
        │   │   │       max_count=255,
        │   │   │       channels_out=8,
        │   │   │       extra=Embedding.Config(
        │   │   │           channels_in=1,
        │   │   │           channels_out=8,
        │   │   │           padding_idx=None,
        │   │   │           device=None,
        │   │   │           dtype=None,
        │   │   │           lookup_dtype=None,
        │   │   │           shard=None,
        │   │   │           depth_index=(),
        │   │   │           init_weight=torch.nn.init.normal_
        │   │   │       )
        │   │   │   ),
        │   │   │   conditions=GraphConditions.Config(
        │   │   │       adduct_groups=[
        │   │   │           [0, 3],
        │   │   │           [0, 2],
        │   │   │           [0, 2],
        │   │   │           [0, 3, 6],
        │   │   │           [0, 3],
        │   │   │           [0],
        │   │   │           [0, 3, 6],
        │   │   │           [0, 3],
        │   │   │           [0, 3],
        │   │   │           [1],
        │   │   │           [1, 4],
        │   │   │           [1, 5],
        │   │   │           [1, 4, 6],
        │   │   │           [1, 4, 7]
        │   │   │       ],
        │   │   │       collision_channels=64,
        │   │   │       collision_scalar=10_000,
        │   │   │       num_instruments=4,
        │   │   │       channels_out=90
        │   │   │   ),
        │   │   │   encoder=MoleculeEncoder.Config(
        │   │   │       channels_in=136,
        │   │   │       channels_out=512,
        │   │   │       projection=Linear.Config(
        │   │   │           channels_in=136,
        │   │   │           channels_out=512,
        │   │   │           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
        │   │   │       ),
        │   │   │       messages=GGNN.Config(
        │   │   │           channels_hidden=512,
        │   │   │           num_layers=6,
        │   │   │           num_edge_types=4,
        │   │   │           edge=Linear.Config(
        │   │   │               channels_in=512,
        │   │   │               channels_out=512,
        │   │   │               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
        │   │   │           ),
        │   │   │           cell=RecurrentCell.Config(
        │   │   │               channels_in=512,
        │   │   │               channels_out=512,
        │   │   │               bias=True
        │   │   │           ),
        │   │   │           init_edge=functools.partial(torch.nn.init.xavier_normal_, gain=1.4142135623730951)
        │   │   │       )
        │   │   │   ),
        │   │   │   head=MoleculeMLP.Config(
        │   │   │   │   channels_in=1_831,
        │   │   │   │   channels_out=1,
        │   │   │   │   channels_hidden=512,
        │   │   │   │   hidden=Sequential.Config(
        │   │   │   │       channels_in=1_831,
        │   │   │   │       channels_out=512,
        │   │   │   │       elements=[
        │   │   │   │           Linear.Config(
        │   │   │   │               channels_in=1_831,
        │   │   │   │               channels_out=512,
        │   │   │   │               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
        │   │   │   │           ),
        │   │   │   │           Activation.Config(
        │   │   │   │               channels_in=512,
        │   │   │   │               channels_out=512,
        │   │   │   │               act=priml.model.swiglu.relu
        │   │   │   │           ),
        │   │   │   │           Dropout.Config(
        │   │   │   │               channels_in=512,
        │   │   │   │               channels_out=512,
        │   │   │   │               p=0.2,
        │   │   │   │               inplace=False
        │   │   │   │           )
        │   │   │   │       ],
        │   │   │   │       repeat=1,
        │   │   │   │       depth_index=()
        │   │   │   │   ),
        │   │   │   │   discarded=Linear.Config(
        │   │   │   │       channels_in=512,
        │   │   │   │       channels_out=512,
        │   │   │   │       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
        │   │   │   │   ),
        │   │   │   │   output=Linear.Config(
        │   │   │   │       channels_in=512,
        │   │   │   │       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
        │   │   │   │   )
        │   │   │   )
        │   │   ),
        │   │   intensity=GraphIntensity.Config(
        │   │   │   channels_in=32,
        │   │   │   channels_hidden=256,
        │   │   │   num_elements=18,
        │   │   │   max_broken=6,
        │   │   │   formula=FormulaEncoding.Config(
        │   │   │       max_count=255,
        │   │   │       channels_out=8,
        │   │   │       extra=Embedding.Config(
        │   │   │           channels_in=1,
        │   │   │           channels_out=8,
        │   │   │           padding_idx=None,
        │   │   │           device=None,
        │   │   │           dtype=None,
        │   │   │           lookup_dtype=None,
        │   │   │           shard=None,
        │   │   │           depth_index=(),
        │   │   │           init_weight=torch.nn.init.normal_
        │   │   │       )
        │   │   │   ),
        │   │   │   conditions=GraphConditions.Config(
        │   │   │       adduct_groups=[
        │   │   │           [0, 3],
        │   │   │           [0, 2],
        │   │   │           [0, 2],
        │   │   │           [0, 3, 6],
        │   │   │           [0, 3],
        │   │   │           [0],
        │   │   │           [0, 3, 6],
        │   │   │           [0, 3],
        │   │   │           [0, 3],
        │   │   │           [1],
        │   │   │           [1, 4],
        │   │   │           [1, 5],
        │   │   │           [1, 4, 6],
        │   │   │           [1, 4, 7]
        │   │   │       ],
        │   │   │       collision_channels=64,
        │   │   │       collision_scalar=10_000,
        │   │   │       num_instruments=4,
        │   │   │       channels_out=90
        │   │   │   ),
        │   │   │   encoder=MoleculeEncoder.Config(
        │   │   │       channels_in=122,
        │   │   │       channels_out=256,
        │   │   │       projection=Linear.Config(
        │   │   │           channels_in=122,
        │   │   │           channels_out=256,
        │   │   │           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
        │   │   │       ),
        │   │   │       messages=GGNN.Config(
        │   │   │           channels_hidden=256,
        │   │   │           num_layers=4,
        │   │   │           num_edge_types=4,
        │   │   │           edge=Linear.Config(
        │   │   │               channels_in=256,
        │   │   │               channels_out=256,
        │   │   │               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
        │   │   │           ),
        │   │   │           cell=RecurrentCell.Config(
        │   │   │               channels_in=256,
        │   │   │               channels_out=256,
        │   │   │               bias=True
        │   │   │           ),
        │   │   │           init_edge=functools.partial(torch.nn.init.xavier_normal_, gain=1.4142135623730951)
        │   │   │       )
        │   │   │   ),
        │   │   │   intermediate=MoleculeMLP.Config(
        │   │   │   │   channels_in=1_063,
        │   │   │   │   channels_out=256,
        │   │   │   │   channels_hidden=256,
        │   │   │   │   hidden=Sequential.Config(
        │   │   │   │       channels_in=1_063,
        │   │   │   │       channels_out=256,
        │   │   │   │       elements=[
        │   │   │   │           Linear.Config(
        │   │   │   │               channels_in=1_063,
        │   │   │   │               channels_out=256,
        │   │   │   │               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
        │   │   │   │           ),
        │   │   │   │           Activation.Config(
        │   │   │   │               channels_in=256,
        │   │   │   │               channels_out=256,
        │   │   │   │               act=priml.model.swiglu.relu
        │   │   │   │           ),
        │   │   │   │           Dropout.Config(
        │   │   │   │               channels_in=256,
        │   │   │   │               channels_out=256,
        │   │   │   │               p=0.1,
        │   │   │   │               inplace=False
        │   │   │   │           )
        │   │   │   │       ],
        │   │   │   │       repeat=1,
        │   │   │   │       depth_index=()
        │   │   │   │   ),
        │   │   │   │   discarded=Linear.Config(
        │   │   │   │       channels_in=256,
        │   │   │   │       channels_out=256,
        │   │   │   │       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
        │   │   │   │   ),
        │   │   │   │   output=Linear.Config(
        │   │   │   │       channels_in=256,
        │   │   │   │       channels_out=256,
        │   │   │   │       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
        │   │   │   │   )
        │   │   │   ),
        │   │   │   num_layers=3,
        │   │   │   block=FragmentBlock.Config(
        │   │   │   │   channels_in=256,
        │   │   │   │   channels_out=256,
        │   │   │   │   attn=FragmentAttention.Config(
        │   │   │   │   │   channels_in=256,
        │   │   │   │   │   channels_out=256,
        │   │   │   │   │   attention=Attention.Config(
        │   │   │   │   │       channels_in=256,
        │   │   │   │   │       channels_out=256,
        │   │   │   │   │       num_heads=8,
        │   │   │   │   │       channels_head=32,
        │   │   │   │   │       num_heads_kv=8,
        │   │   │   │   │       bias=False,
        │   │   │   │   │       dropout=0.1,
        │   │   │   │   │       causal=False,
        │   │   │   │   │       rope=None,
        │   │   │   │   │       rotate=None,
        │   │   │   │   │       norm_qk=None,
        │   │   │   │   │       share_qk_norm=True,
        │   │   │   │   │       norm_out=None,
        │   │   │   │   │       split_qkv_projection=False,
        │   │   │   │   │       proj_qkv=Linear.Config(
        │   │   │   │   │           channels_in=256,
        │   │   │   │   │           channels_out=768,
        │   │   │   │   │           bias=True,
        │   │   │   │   │           depth_index=(),
        │   │   │   │   │           device=None,
        │   │   │   │   │           dtype=None,
        │   │   │   │   │           init_weight=torch.nn.init.zeros_,
        │   │   │   │   │           init_bias=torch.nn.init.zeros_,
        │   │   │   │   │           linear_fn=<built-in function linear>,
        │   │   │   │   │           shard=None
        │   │   │   │   │       ),
        │   │   │   │   │       proj_out=Linear.Config(
        │   │   │   │   │           channels_in=256,
        │   │   │   │   │           channels_out=256,
        │   │   │   │   │           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
        │   │   │   │   │       ),
        │   │   │   │   │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │       depth_index=(),
        │   │   │   │   │       attn_kernel=FragmentKernel.Config()
        │   │   │   │   │   ),
        │   │   │   │   │   drop=Dropout.Config(
        │   │   │   │   │       channels_in=256,
        │   │   │   │   │       channels_out=256,
        │   │   │   │   │       p=0.1,
        │   │   │   │   │       inplace=False
        │   │   │   │   │   )
        │   │   │   │   ),
        │   │   │   │   ffn=Sequential.Config(
        │   │   │   │   │   channels_in=256,
        │   │   │   │   │   channels_out=256,
        │   │   │   │   │   elements=[
        │   │   │   │   │   │   SwiGLU.Config(
        │   │   │   │   │   │       channels_in=256,
        │   │   │   │   │   │       channels_out=256,
        │   │   │   │   │   │       channels_hidden=1_024,
        │   │   │   │   │   │       expansion=4.0,
        │   │   │   │   │   │       round_to=1,
        │   │   │   │   │   │       gate=False,
        │   │   │   │   │   │       bias=False,
        │   │   │   │   │   │       split_gate_projection=False,
        │   │   │   │   │   │       norm=None,
        │   │   │   │   │   │       act=priml.model.swiglu.relu,
        │   │   │   │   │   │       depth_index=(),
        │   │   │   │   │   │       init_weight=priml.model.init.unit_fan_in_uniform,
        │   │   │   │   │   │       init_weight_out=torch.nn.init.zeros_,
        │   │   │   │   │   │       up_proj=Linear.Config(
        │   │   │   │   │   │           channels_in=256,
        │   │   │   │   │   │           channels_out=1_024,
        │   │   │   │   │   │           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=Dropout.Config(
        │   │   │   │   │   │           channels_in=1_024,
        │   │   │   │   │   │           channels_out=1_024,
        │   │   │   │   │   │           p=0.1,
        │   │   │   │   │   │           inplace=False
        │   │   │   │   │   │       ),
        │   │   │   │   │   │       down_proj=Linear.Config(
        │   │   │   │   │   │           channels_in=1_024,
        │   │   │   │   │   │           channels_out=256,
        │   │   │   │   │   │           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
        │   │   │   │   │   │   ),
        │   │   │   │   │   │   Dropout.Config(
        │   │   │   │   │   │       channels_in=256,
        │   │   │   │   │   │       channels_out=256,
        │   │   │   │   │   │       p=0.1,
        │   │   │   │   │   │       inplace=False
        │   │   │   │   │   │   )
        │   │   │   │   │   ],
        │   │   │   │   │   repeat=1,
        │   │   │   │   │   depth_index=()
        │   │   │   │   ),
        │   │   │   │   norm1=LayerNorm.Config(
        │   │   │   │       channels_in=256,
        │   │   │   │       channels_out=256,
        │   │   │   │       eps=1e-05,
        │   │   │   │       elementwise_affine=True,
        │   │   │   │       device=None,
        │   │   │   │       dtype=None
        │   │   │   │   ),
        │   │   │   │   norm2=LayerNorm.Config(
        │   │   │   │       channels_in=256,
        │   │   │   │       channels_out=256,
        │   │   │   │       eps=1e-05,
        │   │   │   │       elementwise_affine=True,
        │   │   │   │       device=None,
        │   │   │   │       dtype=None
        │   │   │   │   ),
        │   │   │   │   prenorm=False,
        │   │   │   │   checkpoint=False,
        │   │   │   │   depth_index=()
        │   │   │   ),
        │   │   │   max_hydrogen_shift=6,
        │   │   │   output=Linear.Config(
        │   │   │       channels_in=256,
        │   │   │       channels_out=26,
        │   │   │       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
        │   │   │   ),
        │   │   │   num_buckets=15_000,
        │   │   │   max_mass=1_500
        │   │   ),
        │   │   generator_checkpoint=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/ahmedberatozer/casmi26-iceberg/ckpt/gen.pt'),
        │   │   intensity_checkpoint=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/ahmedberatozer/casmi26-iceberg/ckpt/inten.pt'),
        │   │   device='cuda',
        │   │   max_nodes=100,
        │   │   max_depth=3,
        │   │   threshold=0.0,
        │   │   reduction=SumReduction.Config(),
        │   │   max_peaks=100
        │   ),
        │   scheduling=IceScheduling.Config(
        │       length_divisors=[(60, 1), (90, 2), (140, 4)],
        │       longest_divisor=8
        │   ),
        │   scoring=QueryScoring.Config(
        │       search=MainSearch.Config(
        │           floor=0.002,
        │           max_peaks=256,
        │           power=1.0,
        │           entropy_weight=True,
        │           tolerance_da=0.01,
        │           tolerance_ppm=20.0
        │       ),
        │       floor=0.001,
        │       max_peaks=512,
        │       over=2.0,
        │       passes=2
        │   ),
        │   budget=WallBudget.Config(
        │       seconds=5400.0,
        │       reserve_evaluation_sec=0,
        │       clock=<built-in function perf_counter>,
        │       receipt_path=None
        │   ),
        │   batch_size=96,
        │   chunk_jobs=3_000,
        │   asset_seals={   PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/ahmedberatozer/casmi26-iceberg/ckpt/gen.pt'): 'deef0451c1870d2125653bb212c9d95c5b5d9bf80156dfa9b22fedda7db49bea',
        │       PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/ahmedberatozer/casmi26-iceberg/ckpt/inten.pt'): 'c49e330e049d3a8b0095026ed694fcf70045c24305354853732aaa5f89d55068'}
        ),
        glacier=SimulatorStage.Config(
        │   base_dir=PosixPath('/opt/scratch/datasets/casmi2026'),
        │   admission_dir=PosixPath('/opt/scratch/datasets/casmi2026/native-chemistry/simulator-admission'),
        │   backend=GlacierBackend.Config(
        │   │   base_dir=PosixPath('/opt/scratch/datasets/casmi2026'),
        │   │   root_dir=PosixPath('/opt/scratch/datasets/casmi2026/native-chemistry/glacier-roots'),
        │   │   model=Glacier.Config(
        │   │   │   channels_hidden=256,
        │   │   │   heads=8,
        │   │   │   channels_in=88,
        │   │   │   intensity_decoder_layers=4,
        │   │   │   intensity_encoder_layers=3,
        │   │   │   max_broken=6,
        │   │   │   num_elements=18,
        │   │   │   formula=FormulaEncoding.Config(
        │   │   │       max_count=255,
        │   │   │       channels_out=8,
        │   │   │       extra=Embedding.Config(
        │   │   │           channels_in=1,
        │   │   │           channels_out=8,
        │   │   │           padding_idx=None,
        │   │   │           device=None,
        │   │   │           dtype=None,
        │   │   │           lookup_dtype=None,
        │   │   │           shard=None,
        │   │   │           depth_index=(),
        │   │   │           init_weight=torch.nn.init.normal_
        │   │   │       )
        │   │   │   ),
        │   │   │   conditions=GraphConditions.Config(
        │   │   │       adduct_groups=[
        │   │   │           [0, 3],
        │   │   │           [0, 2],
        │   │   │           [0, 2],
        │   │   │           [0, 3, 6],
        │   │   │           [0, 3],
        │   │   │           [0],
        │   │   │           [0, 3, 6],
        │   │   │           [0, 3],
        │   │   │           [0, 3],
        │   │   │           [1],
        │   │   │           [1, 4],
        │   │   │           [1, 5],
        │   │   │           [1, 4, 6],
        │   │   │           [1, 4, 7]
        │   │   │       ],
        │   │   │       collision_channels=64,
        │   │   │       collision_scalar=10_000,
        │   │   │       num_instruments=4,
        │   │   │       channels_out=90
        │   │   │   ),
        │   │   │   root=Graphormer.Config(
        │   │   │   │   channels_in=178,
        │   │   │   │   channels_hidden=256,
        │   │   │   │   heads=8,
        │   │   │   │   num_layers=5,
        │   │   │   │   edge_distances=10,
        │   │   │   │   max_hops=10,
        │   │   │   │   init_std=0.02,
        │   │   │   │   node=Linear.Config(
        │   │   │   │       channels_in=178,
        │   │   │   │       channels_out=256,
        │   │   │   │       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
        │   │   │   │   ),
        │   │   │   │   degree=Embedding.Config(
        │   │   │   │       channels_in=8,
        │   │   │   │       channels_out=256,
        │   │   │   │       padding_idx=0,
        │   │   │   │       device=None,
        │   │   │   │       dtype=None,
        │   │   │   │       lookup_dtype=None,
        │   │   │   │       shard=None,
        │   │   │   │       depth_index=(),
        │   │   │   │       init_weight=torch.nn.init.normal_
        │   │   │   │   ),
        │   │   │   │   token=Embedding.Config(
        │   │   │   │       channels_in=1,
        │   │   │   │       channels_out=256,
        │   │   │   │       padding_idx=None,
        │   │   │   │       device=None,
        │   │   │   │       dtype=None,
        │   │   │   │       lookup_dtype=None,
        │   │   │   │       shard=None,
        │   │   │   │       depth_index=(),
        │   │   │   │       init_weight=torch.nn.init.normal_
        │   │   │   │   ),
        │   │   │   │   edge=Linear.Config(
        │   │   │   │       channels_in=28,
        │   │   │   │       channels_out=8,
        │   │   │   │       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
        │   │   │   │   ),
        │   │   │   │   distance=Embedding.Config(
        │   │   │   │       channels_in=640,
        │   │   │   │       channels_out=1,
        │   │   │   │       padding_idx=None,
        │   │   │   │       device=None,
        │   │   │   │       dtype=None,
        │   │   │   │       lookup_dtype=None,
        │   │   │   │       shard=None,
        │   │   │   │       depth_index=(),
        │   │   │   │       init_weight=torch.nn.init.normal_
        │   │   │   │   ),
        │   │   │   │   spatial=Embedding.Config(
        │   │   │   │       channels_in=1_025,
        │   │   │   │       channels_out=8,
        │   │   │   │       padding_idx=0,
        │   │   │   │       device=None,
        │   │   │   │       dtype=None,
        │   │   │   │       lookup_dtype=None,
        │   │   │   │       shard=None,
        │   │   │   │       depth_index=(),
        │   │   │   │       init_weight=torch.nn.init.normal_
        │   │   │   │   ),
        │   │   │   │   virtual=Embedding.Config(
        │   │   │   │       channels_in=1,
        │   │   │   │       channels_out=8,
        │   │   │   │       padding_idx=None,
        │   │   │   │       device=None,
        │   │   │   │       dtype=None,
        │   │   │   │       lookup_dtype=None,
        │   │   │   │       shard=None,
        │   │   │   │       depth_index=(),
        │   │   │   │       init_weight=torch.nn.init.normal_
        │   │   │   │   ),
        │   │   │   │   block=GraphBlock.Config(
        │   │   │   │   │   channels_in=256,
        │   │   │   │   │   channels_out=256,
        │   │   │   │   │   attn=GraphAttention.Config(
        │   │   │   │   │   │   channels_in=256,
        │   │   │   │   │   │   channels_out=256,
        │   │   │   │   │   │   attention=Attention.Config(
        │   │   │   │   │   │       channels_in=256,
        │   │   │   │   │   │       channels_out=256,
        │   │   │   │   │   │       num_heads=8,
        │   │   │   │   │   │       channels_head=32,
        │   │   │   │   │   │       num_heads_kv=8,
        │   │   │   │   │   │       bias=False,
        │   │   │   │   │   │       dropout=0.1,
        │   │   │   │   │   │       causal=False,
        │   │   │   │   │   │       rope=None,
        │   │   │   │   │   │       rotate=None,
        │   │   │   │   │   │       norm_qk=None,
        │   │   │   │   │   │       share_qk_norm=True,
        │   │   │   │   │   │       norm_out=None,
        │   │   │   │   │   │       split_qkv_projection=False,
        │   │   │   │   │   │       proj_qkv=Linear.Config(
        │   │   │   │   │   │           channels_in=256,
        │   │   │   │   │   │           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
        │   │   │   │   │   │       ),
        │   │   │   │   │   │       proj_out=Linear.Config(
        │   │   │   │   │   │           channels_in=256,
        │   │   │   │   │   │           channels_out=256,
        │   │   │   │   │   │           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
        │   │   │   │   │   │       ),
        │   │   │   │   │   │       init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │   │       depth_index=(),
        │   │   │   │   │   │       attn_kernel=GraphKernel.Config()
        │   │   │   │   │   │   ),
        │   │   │   │   │   │   dropout=Dropout.Config(
        │   │   │   │   │   │       channels_in=256,
        │   │   │   │   │   │       channels_out=256,
        │   │   │   │   │   │       p=0.1,
        │   │   │   │   │   │       inplace=False
        │   │   │   │   │   │   )
        │   │   │   │   │   ),
        │   │   │   │   │   ffn=Sequential.Config(
        │   │   │   │   │   │   channels_in=256,
        │   │   │   │   │   │   channels_out=256,
        │   │   │   │   │   │   elements=[
        │   │   │   │   │   │   │   SwiGLU.Config(
        │   │   │   │   │   │   │       channels_in=256,
        │   │   │   │   │   │   │       channels_out=256,
        │   │   │   │   │   │   │       channels_hidden=1_024,
        │   │   │   │   │   │   │       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=256,
        │   │   │   │   │   │   │           channels_out=1_024,
        │   │   │   │   │   │   │           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=Dropout.Config(
        │   │   │   │   │   │   │           channels_in=1_024,
        │   │   │   │   │   │   │           channels_out=1_024,
        │   │   │   │   │   │   │           p=0.1,
        │   │   │   │   │   │   │           inplace=False
        │   │   │   │   │   │   │       ),
        │   │   │   │   │   │   │       down_proj=Linear.Config(
        │   │   │   │   │   │   │           channels_in=1_024,
        │   │   │   │   │   │   │           channels_out=256,
        │   │   │   │   │   │   │           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
        │   │   │   │   │   │   │   ),
        │   │   │   │   │   │   │   Dropout.Config(
        │   │   │   │   │   │   │       channels_in=256,
        │   │   │   │   │   │   │       channels_out=256,
        │   │   │   │   │   │   │       p=0.1,
        │   │   │   │   │   │   │       inplace=False
        │   │   │   │   │   │   │   )
        │   │   │   │   │   │   ],
        │   │   │   │   │   │   repeat=1,
        │   │   │   │   │   │   depth_index=()
        │   │   │   │   │   ),
        │   │   │   │   │   norm1=LayerNorm.Config(
        │   │   │   │   │       channels_in=256,
        │   │   │   │   │       channels_out=256,
        │   │   │   │   │       eps=1e-05,
        │   │   │   │   │       elementwise_affine=True,
        │   │   │   │   │       device=None,
        │   │   │   │   │       dtype=None
        │   │   │   │   │   ),
        │   │   │   │   │   norm2=LayerNorm.Config(
        │   │   │   │   │       channels_in=256,
        │   │   │   │   │       channels_out=256,
        │   │   │   │   │       eps=1e-05,
        │   │   │   │   │       elementwise_affine=True,
        │   │   │   │   │       device=None,
        │   │   │   │   │       dtype=None
        │   │   │   │   │   ),
        │   │   │   │   │   prenorm=False,
        │   │   │   │   │   checkpoint=False,
        │   │   │   │   │   depth_index=()
        │   │   │   │   ),
        │   │   │   │   dropout=Dropout.Config(
        │   │   │   │       channels_in=-1,
        │   │   │   │       channels_out=-1,
        │   │   │   │       p=0.1,
        │   │   │   │       inplace=False
        │   │   │   │   )
        │   │   │   ),
        │   │   │   formula_mapper=Linear.Config(
        │   │   │       channels_in=400,
        │   │   │       channels_out=256,
        │   │   │       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
        │   │   │   ),
        │   │   │   token_mapper=Linear.Config(
        │   │   │       channels_in=551,
        │   │   │       channels_out=256,
        │   │   │       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
        │   │   │   ),
        │   │   │   queries=Slots.Config(
        │   │   │   │   count=200,
        │   │   │   │   num_layers=4,
        │   │   │   │   embedding=Embedding.Config(
        │   │   │   │       channels_in=200,
        │   │   │   │       channels_out=256,
        │   │   │   │       padding_idx=None,
        │   │   │   │       device=None,
        │   │   │   │       dtype=None,
        │   │   │   │       lookup_dtype=None,
        │   │   │   │       shard=None,
        │   │   │   │       depth_index=(),
        │   │   │   │       init_weight=functools.partial(torch.nn.init.orthogonal_, gain=0.02)
        │   │   │   │   ),
        │   │   │   │   decoder=Decoder.Config(
        │   │   │   │   │   channels_in=256,
        │   │   │   │   │   attention=PackedAttention.Config(
        │   │   │   │   │       projections=AttentionProjections.Config(
        │   │   │   │   │           channels_in=256,
        │   │   │   │   │           channels_out=256,
        │   │   │   │   │           num_heads=8,
        │   │   │   │   │           channels_head=32,
        │   │   │   │   │           num_heads_kv=8,
        │   │   │   │   │           bias=False,
        │   │   │   │   │           dropout=0.1,
        │   │   │   │   │           causal=False,
        │   │   │   │   │           rope=None,
        │   │   │   │   │           rotate=None,
        │   │   │   │   │           norm_qk=None,
        │   │   │   │   │           share_qk_norm=True,
        │   │   │   │   │           norm_out=None,
        │   │   │   │   │           split_qkv_projection=False,
        │   │   │   │   │           proj_qkv=Linear.Config(
        │   │   │   │   │               channels_in=256,
        │   │   │   │   │               channels_out=768,
        │   │   │   │   │               bias=True,
        │   │   │   │   │               depth_index=(),
        │   │   │   │   │               device=None,
        │   │   │   │   │               dtype=None,
        │   │   │   │   │               init_weight=torch.nn.init.zeros_,
        │   │   │   │   │               init_bias=torch.nn.init.zeros_,
        │   │   │   │   │               linear_fn=<built-in function linear>,
        │   │   │   │   │               shard=None
        │   │   │   │   │           ),
        │   │   │   │   │           proj_out=Linear.Config(
        │   │   │   │   │               channels_in=256,
        │   │   │   │   │               channels_out=256,
        │   │   │   │   │               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
        │   │   │   │   │           ),
        │   │   │   │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │   │           depth_index=()
        │   │   │   │   │       )
        │   │   │   │   │   ),
        │   │   │   │   │   expand=Linear.Config(
        │   │   │   │   │       channels_in=256,
        │   │   │   │   │       channels_out=1_024,
        │   │   │   │   │       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
        │   │   │   │   │   ),
        │   │   │   │   │   contract=Linear.Config(
        │   │   │   │   │       channels_in=1_024,
        │   │   │   │   │       channels_out=256,
        │   │   │   │   │       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
        │   │   │   │   │   ),
        │   │   │   │   │   norm=LayerNorm.Config(
        │   │   │   │   │       channels_in=256,
        │   │   │   │   │       channels_out=256,
        │   │   │   │   │       eps=1e-05,
        │   │   │   │   │       elementwise_affine=True,
        │   │   │   │   │       device=None,
        │   │   │   │   │       dtype=None
        │   │   │   │   │   ),
        │   │   │   │   │   dropout=Dropout.Config(
        │   │   │   │   │       channels_in=-1,
        │   │   │   │   │       channels_out=-1,
        │   │   │   │   │       p=0.1,
        │   │   │   │   │       inplace=False
        │   │   │   │   │   ),
        │   │   │   │   │   activation=Activation.Config(
        │   │   │   │   │       channels_in=-1,
        │   │   │   │   │       channels_out=-1,
        │   │   │   │   │       act=priml.model.swiglu.relu
        │   │   │   │   │   )
        │   │   │   │   )
        │   │   │   ),
        │   │   │   cardinality=Linear.Config(
        │   │   │       channels_in=256,
        │   │   │       channels_out=4,
        │   │   │       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
        │   │   │   ),
        │   │   │   decoder=Decoder.Config(
        │   │   │   │   channels_in=256,
        │   │   │   │   attention=PackedAttention.Config(
        │   │   │   │       projections=AttentionProjections.Config(
        │   │   │   │           channels_in=256,
        │   │   │   │           channels_out=256,
        │   │   │   │           num_heads=8,
        │   │   │   │           channels_head=32,
        │   │   │   │           num_heads_kv=8,
        │   │   │   │           bias=False,
        │   │   │   │           dropout=0.1,
        │   │   │   │           causal=False,
        │   │   │   │           rope=None,
        │   │   │   │           rotate=None,
        │   │   │   │           norm_qk=None,
        │   │   │   │           share_qk_norm=True,
        │   │   │   │           norm_out=None,
        │   │   │   │           split_qkv_projection=False,
        │   │   │   │           proj_qkv=Linear.Config(
        │   │   │   │               channels_in=256,
        │   │   │   │               channels_out=768,
        │   │   │   │               bias=True,
        │   │   │   │               depth_index=(),
        │   │   │   │               device=None,
        │   │   │   │               dtype=None,
        │   │   │   │               init_weight=torch.nn.init.zeros_,
        │   │   │   │               init_bias=torch.nn.init.zeros_,
        │   │   │   │               linear_fn=<built-in function linear>,
        │   │   │   │               shard=None
        │   │   │   │           ),
        │   │   │   │           proj_out=Linear.Config(
        │   │   │   │               channels_in=256,
        │   │   │   │               channels_out=256,
        │   │   │   │               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
        │   │   │   │           ),
        │   │   │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │           depth_index=()
        │   │   │   │       )
        │   │   │   │   ),
        │   │   │   │   expand=Linear.Config(
        │   │   │   │       channels_in=256,
        │   │   │   │       channels_out=1_024,
        │   │   │   │       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
        │   │   │   │   ),
        │   │   │   │   contract=Linear.Config(
        │   │   │   │       channels_in=1_024,
        │   │   │   │       channels_out=256,
        │   │   │   │       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
        │   │   │   │   ),
        │   │   │   │   norm=LayerNorm.Config(
        │   │   │   │       channels_in=256,
        │   │   │   │       channels_out=256,
        │   │   │   │       eps=1e-05,
        │   │   │   │       elementwise_affine=True,
        │   │   │   │       device=None,
        │   │   │   │       dtype=None
        │   │   │   │   ),
        │   │   │   │   dropout=Dropout.Config(
        │   │   │   │       channels_in=-1,
        │   │   │   │       channels_out=-1,
        │   │   │   │       p=0.1,
        │   │   │   │       inplace=False
        │   │   │   │   ),
        │   │   │   │   activation=Activation.Config(
        │   │   │   │       channels_in=-1,
        │   │   │   │       channels_out=-1,
        │   │   │   │       act=priml.model.swiglu.relu
        │   │   │   │   )
        │   │   │   ),
        │   │   │   encoder=Encoder.Config(
        │   │   │   │   channels_in=256,
        │   │   │   │   attention=PackedAttention.Config(
        │   │   │   │       projections=AttentionProjections.Config(
        │   │   │   │           channels_in=256,
        │   │   │   │           channels_out=256,
        │   │   │   │           num_heads=8,
        │   │   │   │           channels_head=32,
        │   │   │   │           num_heads_kv=8,
        │   │   │   │           bias=False,
        │   │   │   │           dropout=0.1,
        │   │   │   │           causal=False,
        │   │   │   │           rope=None,
        │   │   │   │           rotate=None,
        │   │   │   │           norm_qk=None,
        │   │   │   │           share_qk_norm=True,
        │   │   │   │           norm_out=None,
        │   │   │   │           split_qkv_projection=False,
        │   │   │   │           proj_qkv=Linear.Config(
        │   │   │   │               channels_in=256,
        │   │   │   │               channels_out=768,
        │   │   │   │               bias=True,
        │   │   │   │               depth_index=(),
        │   │   │   │               device=None,
        │   │   │   │               dtype=None,
        │   │   │   │               init_weight=torch.nn.init.zeros_,
        │   │   │   │               init_bias=torch.nn.init.zeros_,
        │   │   │   │               linear_fn=<built-in function linear>,
        │   │   │   │               shard=None
        │   │   │   │           ),
        │   │   │   │           proj_out=Linear.Config(
        │   │   │   │               channels_in=256,
        │   │   │   │               channels_out=256,
        │   │   │   │               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
        │   │   │   │           ),
        │   │   │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │           depth_index=()
        │   │   │   │       )
        │   │   │   │   ),
        │   │   │   │   expand=Linear.Config(
        │   │   │   │       channels_in=256,
        │   │   │   │       channels_out=1_024,
        │   │   │   │       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
        │   │   │   │   ),
        │   │   │   │   contract=Linear.Config(
        │   │   │   │       channels_in=1_024,
        │   │   │   │       channels_out=256,
        │   │   │   │       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
        │   │   │   │   ),
        │   │   │   │   norm=LayerNorm.Config(
        │   │   │   │       channels_in=256,
        │   │   │   │       channels_out=256,
        │   │   │   │       eps=1e-05,
        │   │   │   │       elementwise_affine=True,
        │   │   │   │       device=None,
        │   │   │   │       dtype=None
        │   │   │   │   ),
        │   │   │   │   dropout=Dropout.Config(
        │   │   │   │       channels_in=-1,
        │   │   │   │       channels_out=-1,
        │   │   │   │       p=0.1,
        │   │   │   │       inplace=False
        │   │   │   │   ),
        │   │   │   │   activation=Activation.Config(
        │   │   │   │       channels_in=-1,
        │   │   │   │       channels_out=-1,
        │   │   │   │       act=priml.model.swiglu.relu
        │   │   │   │   )
        │   │   │   ),
        │   │   │   output=Linear.Config(
        │   │   │       channels_in=256,
        │   │   │       channels_out=26,
        │   │   │       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
        │   │   │   ),
        │   │   │   breakpoints=Breakpoints.Config(tau=15)
        │   │   ),
        │   │   checkpoint=PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/ahmedberatozer/casmi26-glacier/ckpt/glacier.pt'),
        │   │   device='cuda',
        │   │   walk_steps=14,
        │   │   reduction=MaxReduction.Config(),
        │   │   max_peaks=100
        │   ),
        │   scheduling=GlacierScheduling.Config(
        │       batch_atoms_squared=1_600,
        │       fail_fast_count=64
        │   ),
        │   scoring=QueryScoring.Config(
        │       search=MainSearch.Config(
        │           floor=0.002,
        │           max_peaks=256,
        │           power=1.0,
        │           entropy_weight=True,
        │           tolerance_da=0.01,
        │           tolerance_ppm=20.0
        │       ),
        │       floor=0.001,
        │       max_peaks=512,
        │       over=2.0,
        │       passes=2
        │   ),
        │   budget=WallBudget.Config(
        │       seconds=4000.0,
        │       reserve_evaluation_sec=0,
        │       clock=<built-in function perf_counter>,
        │       receipt_path=None
        │   ),
        │   batch_size=32,
        │   chunk_jobs=3_000,
        │   asset_seals={   PosixPath('/opt/scratch/datasets/casmi2026/inputs/datasets/ahmedberatozer/casmi26-glacier/ckpt/glacier.pt'): '9192e4c2eeafe83317d5a2c2641e1fc4d68934488757a4faafeb66f2ab397a12'}
        ),
        identities=IdentityTable.Config(
            base_dir=PosixPath('/opt/scratch/datasets/casmi2026'),
            working_dir=PosixPath('/opt/scratch/datasets/casmi2026/native-chemistry/identities')
        ),
        graphs=GapGraphTable.Config(
            base_dir=PosixPath('/opt/scratch/datasets/casmi2026'),
            working_dir=PosixPath('/opt/scratch/datasets/casmi2026/native-chemistry/fragment-gap')
        ),
        fragment_gap=FragmentGap.Config(
            level_weights=[1.0, 0.6],
            max_bonds=[100],
            intensity_power=1.0,
            mz_power=0.0,
            tol=0.005,
            h_shifts=[-2, -1, 0, 1, 2, 3],
            max_spectra=16
        ),
        policy=FusionPolicy.Config(
            legacy=RankFusion.Config(alpha=0.6, offset=3.0, window=40),
            control=RankFusion.Config(alpha=0.6, offset=3.0, window=40),
            final=RankFusion.Config(alpha=0.6, offset=3.0, window=40),
            fallback=SimulatorFallback.Config(gap_weight=0.5),
            length=25,
            min_covered=2,
            tie_eps=1e-09,
            ice_weight=1.0,
            glacier_weight=1.0
        ),
        budget=ListBudget.Config(
            main_candidates=60,
            second_candidates=40,
            pubchem_candidates=25,
            legacy_rerank_top_n=60,
            pubchem_rerank_top_n=25,
            main_limit=25,
            insertion_limit=25
        ),
        gate=PubChemGate.Config(
            library_gate=0.9,
            margin_aggressive=600.0,
            margin_promote=250.0,
            outlier_promote=6.0,
            popularity_promote=5.0,
            slots_aggressive=[2, 4, 6, 8, 10],
            slots_gentle=[4, 8, 12, 16, 20],
            limit=25
        ),
        provenance=PublicationProvenance.Config(
            control_experiment='E004',
            final_experiment='E006',
            recipe_sha256='81ebe7b6972bdebdf3d1227d77565e35a71b4b3418b155e3035ccaaff9d53005',
            notebook_sha256='1f0f9b3026a93ddb5ad6fa6993c8ac0a4bfaddc54594eb1dd6dda7a20d8ef0d6',
            source_code_sha256='eee496a371fce179b4aeda1c5082a58212c724eaf1ca4ba0c943bbe8d7196ef2',
            helper_source_sha256='a027a69ddda2c2e14639e3ec5196d5f1eb2ab57cba10ef929723148f31fc8f34',
            metric_source_sha256='c61d163f1bb768313fd0863822952f035f343924f37e3f5e4b5883ccef7e6201',
            source_version_id=356_740_494,
            source_version=80,
            ice_policy_source_version_id=356_785_434,
            ice_policy_source='samanyu1808/casmi26-iceberg-where-glacier-is-blind-lb-0-438/1'
        ),
        device='cuda',
        clock=<built-in function perf_counter>
    )
