CharmTrainLoop.Config(
        study_name='arcagi',
        experiment_name='exp010',
        base_dir='/opt/scratch',
        working_dir=PosixPath('/opt/scratch/runs/arcagi/exp010'),
        doc='',
        step=PooledActTrainStep.Config(
        │   model=SingleLatentGridSolver.Config(
        │   │   channels_in=512,
        │   │   num_layers=4,
        │   │   embedding=GridEmbedding.Config(
        │   │       channels_in=12,
        │   │       channels_out=512,
        │   │       grid_shape=(30, 30),
        │   │       channels=[],
        │   │       init_weight=functools.partial(priml.model.init.fan_in_truncated_normal, variance_correction=True),
        │   │       token_table=Embedding.Config(
        │   │           channels_in=12,
        │   │           channels_out=512,
        │   │           padding_idx=None,
        │   │           device=None,
        │   │           dtype=None,
        │   │           lookup_dtype=torch.bfloat16,
        │   │           shard=None,
        │   │           depth_index=(),
        │   │           init_weight=functools.partial(priml.model.init.fan_in_truncated_normal, variance_correction=True)
        │   │       ),
        │   │       embed_scale=1.0
        │   │   ),
        │   │   block=TransformerBlock.Config(
        │   │   │   channels_in=512,
        │   │   │   channels_out=512,
        │   │   │   attn=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=False,
        │   │   │           depth_index=(),
        │   │   │           device=None,
        │   │   │           dtype=None,
        │   │   │           init_weight=functools.partial(priml.model.init.fan_in_truncated_normal, variance_correction=True),
        │   │   │           init_bias=torch.nn.init.zeros_,
        │   │   │           linear_fn=priml.math.linear.cast_parameters_to_input,
        │   │   │           shard=None
        │   │   │       ),
        │   │   │       proj_out=Linear.Config(
        │   │   │           channels_in=512,
        │   │   │           channels_out=512,
        │   │   │           bias=False,
        │   │   │           depth_index=(),
        │   │   │           device=None,
        │   │   │           dtype=None,
        │   │   │           init_weight=functools.partial(priml.model.init.fan_in_truncated_normal, variance_correction=True),
        │   │   │           init_bias=torch.nn.init.zeros_,
        │   │   │           linear_fn=priml.math.linear.cast_parameters_to_input,
        │   │   │           shard=None
        │   │   │       ),
        │   │   │       init_weight=functools.partial(priml.model.init.fan_in_truncated_normal, variance_correction=True),
        │   │   │       depth_index=(),
        │   │   │       attn_kernel=SdpaFused.Config()
        │   │   │   ),
        │   │   │   ffn=SwiGLU.Config(
        │   │   │   │   channels_in=512,
        │   │   │   │   channels_out=512,
        │   │   │   │   channels_hidden=1_536,
        │   │   │   │   expansion=4.0,
        │   │   │   │   round_to=256,
        │   │   │   │   gate=True,
        │   │   │   │   bias=False,
        │   │   │   │   split_gate_projection=False,
        │   │   │   │   norm=None,
        │   │   │   │   act=priml.model.swiglu.silu,
        │   │   │   │   depth_index=(),
        │   │   │   │   init_weight=functools.partial(priml.model.init.fan_in_truncated_normal, variance_correction=True),
        │   │   │   │   init_weight_out=functools.partial(priml.model.init.fan_in_truncated_normal, variance_correction=True),
        │   │   │   │   up_proj=Linear.Config(
        │   │   │   │       channels_in=512,
        │   │   │   │       channels_out=3_072,
        │   │   │   │       bias=False,
        │   │   │   │       depth_index=(),
        │   │   │   │       device=None,
        │   │   │   │       dtype=None,
        │   │   │   │       init_weight=functools.partial(priml.model.init.fan_in_truncated_normal, variance_correction=True),
        │   │   │   │       init_bias=torch.nn.init.zeros_,
        │   │   │   │       linear_fn=priml.math.linear.cast_parameters_to_input,
        │   │   │   │       shard=None
        │   │   │   │   ),
        │   │   │   │   middle=SequenceConv1d.Config(
        │   │   │   │       channels_in=1_536,
        │   │   │   │       channels_out=1_536,
        │   │   │   │       conv=Conv1d.Config(
        │   │   │   │           channels_in=1_536,
        │   │   │   │           channels_out=1_536,
        │   │   │   │           kernel_size=2,
        │   │   │   │           stride=1,
        │   │   │   │           padding=1,
        │   │   │   │           dilation=1,
        │   │   │   │           groups=1_536,
        │   │   │   │           bias=True,
        │   │   │   │           depth_index=(),
        │   │   │   │           device=None,
        │   │   │   │           dtype=torch.bfloat16,
        │   │   │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │   │   │           init_bias=functools.partial(torch.nn.init.uniform_, a=-0.7071067811865476, b=0.7071067811865476)
        │   │   │   │       ),
        │   │   │   │       act=torch.nn.functional.silu
        │   │   │   │   ),
        │   │   │   │   down_proj=Linear.Config(
        │   │   │   │       channels_in=1_536,
        │   │   │   │       channels_out=512,
        │   │   │   │       bias=False,
        │   │   │   │       depth_index=(),
        │   │   │   │       device=None,
        │   │   │   │       dtype=None,
        │   │   │   │       init_weight=functools.partial(priml.model.init.fan_in_truncated_normal, variance_correction=True),
        │   │   │   │       init_bias=torch.nn.init.zeros_,
        │   │   │   │       linear_fn=priml.math.linear.cast_parameters_to_input,
        │   │   │   │       shard=None
        │   │   │   │   ),
        │   │   │   │   shard='colwise'
        │   │   │   ),
        │   │   │   norm1=RMSNorm.Config(
        │   │   │       channels_in=512,
        │   │   │       channels_out=512,
        │   │   │       eps=1e-05,
        │   │   │       elementwise_affine=False,
        │   │   │       device=None,
        │   │   │       dtype=None
        │   │   │   ),
        │   │   │   norm2=RMSNorm.Config(
        │   │   │       channels_in=512,
        │   │   │       channels_out=512,
        │   │   │       eps=1e-05,
        │   │   │       elementwise_affine=False,
        │   │   │       device=None,
        │   │   │       dtype=None
        │   │   │   ),
        │   │   │   prenorm=False,
        │   │   │   checkpoint=False,
        │   │   │   depth_index=()
        │   │   ),
        │   │   recurrence=SingleLatentRecurrence.Config(
        │   │       slow_cycles=1,
        │   │       fast_cycles=12,
        │   │       inner_grad_loops=6
        │   │   ),
        │   │   prefix=ComposedSparsePrefix.Config(
        │   │       channels_out=512,
        │   │       num_tokens=16,
        │   │       forward_dtype=torch.bfloat16,
        │   │       composition=CategorySlotFiLM.Config(
        │   │           num_categories=1_280,
        │   │           num_augmentations=8,
        │   │           num_slots=9,
        │   │           channels_out=512
        │   │       ),
        │   │       residual=SparsePuzzleEmbedding.Config(
        │   │           channels_in=-1,
        │   │           channels_out=32,
        │   │           num_puzzles=1_289_151,
        │   │           num_tokens=1,
        │   │           batch_size=96,
        │   │           init_std=0.0,
        │   │           embed_scale=1.0,
        │   │           dtype=torch.bfloat16,
        │   │           dtype_scale=torch.bfloat16
        │   │       ),
        │   │       residual_projection=Linear.Config(
        │   │           channels_in=32,
        │   │           channels_out=512,
        │   │           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
        │   │       ),
        │   │       residual_gate=None
        │   │   ),
        │   │   num_prefix_tokens=16,
        │   │   vocab_size=12,
        │   │   head=Linear.Config(
        │   │       channels_in=512,
        │   │       channels_out=12,
        │   │       bias=False,
        │   │       depth_index=(),
        │   │       device=None,
        │   │       dtype=None,
        │   │       init_weight=functools.partial(priml.model.init.fan_in_truncated_normal, variance_correction=True),
        │   │       init_bias=torch.nn.init.zeros_,
        │   │       linear_fn=priml.math.linear.cast_parameters_to_input,
        │   │       shard=None
        │   │   ),
        │   │   init_latent=functools.partial(priml.model.init.truncated_normal, std=1.0, variance_correction=True),
        │   │   halt_outputs=1,
        │   │   halt_init_bias=-5.0,
        │   │   dtype=None,
        │   │   rope=RoPE.Config(
        │   │       channels_head=64,
        │   │       reduction_mode='cat',
        │   │       frequencies=HuggingFaceFrequencies.Config(base=10000.0),
        │   │       dtype=torch.float32
        │   │   ),
        │   │   rope_grid_shape=(900,),
        │   │   compile_core=CoreCompile.Config(
        │   │       unit='core',
        │   │       mode='default',
        │   │       fullgraph=False
        │   │   ),
        │   │   dtype_latent_init=None,
        │   │   checkpoint_core=False,
        │   │   block_checkpoint_fraction=0.0,
        │   │   halt_head=Linear.Config(
        │   │       channels_in=512,
        │   │       channels_out=1,
        │   │       bias=True,
        │   │       depth_index=(),
        │   │       device=None,
        │   │       dtype=None,
        │   │       init_weight=functools.partial(priml.model.init.fan_in_truncated_normal, variance_correction=True),
        │   │       init_bias=torch.nn.init.zeros_,
        │   │       linear_fn=priml.math.linear.cast_parameters_to_input,
        │   │       shard=None
        │   │   ),
        │   │   forward_dtype=torch.bfloat16
        │   ),
        │   optimizer=CompositeOptimizer.Config(
        │       optimizers=[
        │           Muon.Config(
        │               lr=0.0001,
        │               weight_decay=0.1,
        │               momentum=0.95,
        │               nesterov=True,
        │               ns_coefficients=(3.4445, -4.775, 2.0315),
        │               eps=1e-07,
        │               ns_steps=5,
        │               reference_numerics=False,
        │               adjust_lr_fn=priml.optimizers.muon.adjust_lr_match_rms_adamw,
        │               momentum_update_fn=priml.optimizers.muon.momentum_add_update,
        │               matrix_signum_fn=functools.partial(priml.math.numeric.matrix_signum_via_coefficients, coefficients=((7.2086, -15.5131, 9.0178), (3.9623, -2.5813, 0.4542), (3.9466, -2.5765, 0.4544), (3.8991, -2.5671, 0.4566), (3.7186, -2.5308, 0.4653), (3.139, -2.3073, 0.4733), (2.1715, -1.5246, 0.3885), (1.8648, -1.2224, 0.3577)), steps=5, eps=1e-07, compute_dtype=torch.float32),
        │               ensemble_dims=0
        │           ),
        │           AuxAdam.Config(
        │               lr=0.0001,
        │               betas=(0.9, 0.95),
        │               eps=1e-08,
        │               weight_decay=0.1
        │           )
        │       ],
        │       select=[
        │           with_ndim(Muon.eligible_tensor, 2),
        │           complement(with_ndim(Muon.eligible_tensor, 2))
        │       ],
        │       require_total=True,
        │       drop_empty=True
        │   ),
        │   loss=SimpleLoss.Config(
        │       loss_fn=priml.loss.simple_loss.bce_with_logits,
        │       target_key='label',
        │       kwargs={'reduction': 'none'},
        │       channels_out=-1
        │   ),
        │   accumulate_grad_batches=1,
        │   drop_partial_accumulation_on_epoch_end=True,
        │   train_budget_steps=inf,
        │   train_budget_sec=inf,
        │   train_budget_epochs=inf,
        │   lr_schedule=PartialConfig(priml.math.schedules.constant),
        │   parallelism=ManualDataParallel.Config(mesh_dim='dp'),
        │   model_quantization=NoModelQuantization.Config(),
        │   activation_memoization=DefaultActivationStorage.Config(),
        │   compile=None,
        │   ema=EMA.Config(
        │       decay=0.999,
        │       update_after_step=0,
        │       update_every=1,
        │       track_buffers=False,
        │       shadow_kind='param_dict',
        │       warmup_seed=False,
        │       decay_schedule=priml.train.ema.constant_decay,
        │       select=priml.optimizers.parameter_filter.trainable,
        │       update_fn=priml.baselines.arcagi.experiments._source_order_ema_update,
        │       buffer_update_fn=priml.train.ema.copy_buffers
        │   ),
        │   gradient_clip_norm=inf,
        │   skip_step_on_nonfinite_grad=False,
        │   device_init='eager',
        │   dtype_autocast=None,
        │   autocast_cache_enabled=False,
        │   sparse_optimizer=SignSGD.Config(
        │       lr=0.01,
        │       weight_decay=0.1,
        │       aggregate_distributed=True
        │   ),
        │   optimizer_lr_policy=priml.baselines.arcagi.train_step.apply_direct_schedule_to_body_and_sparse_optimizer_lrs,
        │   pool=AtomicPool.Config(
        │       batch_size=96,
        │       max_steps=16,
        │       halting=HaltTraining.Config(
        │           weight=0.5,
        │           exploration=GlobalSampledMinimum.Config(prob=0.1),
        │           seed=0
        │       ),
        │       feedback=None,
        │       start=LearnedStart.Config(),
        │       grid_len=900,
        │       seq_len=916,
        │       channels_hidden=512,
        │       steps_dtype=torch.int32,
        │       dtype=torch.bfloat16
        │   ),
        │   pool_transition=BeforeBackwardPoolTransition.Config(),
        │   ema_initialization=priml.train.ema.EMA.initialize,
        │   gradient_reduction=SumGradients.Config(mesh_dim='dp'),
        │   metric_reporter=PartialConfig(
        │       priml.baselines.arcagi.experiments._report_source_act_metrics,
        │       global_batch_size=768
        │   ),
        │   token_loss=StablemaxTokens.Config(
        │       log_normalizer=priml.math.loss.log_stablemax
        │   ),
        │   objective=ACTLoss.Config(
        │       token_loss=StablemaxTokens.Config(
        │           log_normalizer=priml.math.loss.log_stablemax_via_probabilities
        │       ),
        │       halt_loss=SimpleLoss.Config(
        │           loss_fn=priml.loss.simple_loss.bce_with_logits,
        │           target_key='halt_target',
        │           kwargs={'reduction': 'sum'},
        │           channels_out=-1
        │       ),
        │       ignore_index=-100,
        │       halt_weight=0.5
        │   ),
        │   reduction=MeanOverBatch.Config(batch_size=96),
        │   signals=None,
        │   ignore_label_id=-100,
        │   total_train_steps=2_380_000,
        │   backward_loss_scale=0.0013020833333333333,
        │   lr_step_offset=1,
        │   warmup_steps=2_000,
        │   lr_min_ratio=1.0,
        │   grad_clip_foreach=True,
        │   emulate_precision_casts=False,
        │   log_body_norms=True,
        │   norm_log_interval=100
        ),
        dataset=ArcParquetData.Config(
            base_dir='/opt/scratch',
            train_sources=[
                PairSource.Config(
                    path=PosixPath('/opt/scratch/datasets/charm-release-937316c8/data/augmented/v2/arc2_training_aug1000.parquet'),
                    pair_selection='both',
                    multiplicity=1
                ),
                PairSource.Config(
                    path=PosixPath('/opt/scratch/datasets/charm-release-937316c8/data/augmented/v2/rearc2_aug100.parquet'),
                    pair_selection='both',
                    multiplicity=1
                ),
                PairSource.Config(
                    path=PosixPath('/opt/scratch/datasets/charm-release-937316c8/data/augmented/v2/concept_aug1000.parquet'),
                    pair_selection='both',
                    multiplicity=2
                ),
                PairSource.Config(
                    path=PosixPath('/opt/scratch/datasets/charm-release-937316c8/data/augmented/v2/arc2_eval_aug1000.parquet'),
                    pair_selection='train',
                    multiplicity=2
                )
            ],
            eval_source=PosixPath('/opt/scratch/datasets/charm-release-937316c8/data/augmented/v2/arc2_eval_aug1000.parquet'),
            canvas_size=30,
            batch_size=96,
            eval_batch_size=96,
            num_workers=4,
            seed=0,
            p_translate=0.8,
            rank=-1,
            num_replicas=-1,
            num_tasks=1_280,
            num_instances=1_289_150
        ),
        metrics_train={},
        metrics_eval={   '': ParquetPassK.Config(
                base_dir='/opt/scratch',
                eval_source=PosixPath('/opt/scratch/datasets/charm-release-937316c8/data/augmented/v2/arc2_eval_aug1000.parquet'),
                canvas_size=30,
                pass_ks=(1, 2, 5, 10, 100, 1_000),
                recent_window_size=None,
                aggregated_voting=True
            )},
        checkpointer=Checkpointer.Config(
            base_dir=PosixPath('/opt/scratch/runs/arcagi/exp010'),
            working_dir=PosixPath('/opt/scratch/runs/arcagi/exp010/checkpoints'),
            filename='step_{step:08d}.pt',
            save_every=10_000,
            keep_last_n=-1,
            keep_every=0,
            storer=SyncLocalStateDictStorer.Config(),
            resume=True,
            resume_step=-1,
            allow_checkpoint_overwrite=False,
            best_metric='',
            best_mode='max'
        ),
        profiler=None,
        phase_timer=PhaseTimer.Config(
            enabled=False,
            torch_profile=False,
            base_dir=PosixPath('/opt/scratch/runs/arcagi/exp010'),
            working_dir=PosixPath('/opt/scratch/runs/arcagi/exp010/profiling'),
            heartbeat_interval_sec=20.0,
            fault_dump_interval_sec=40.0,
            cuda_events=False
        ),
        tracker=TrackerList.Config(
            trackers={   'wandb': WandbTracker.Config(
                    project='charm',
                    name='arc-agi-2-charm-rearc2-100',
                    run_id='',
                    group='',
                    mode='online',
                    base_dir=PosixPath('/opt/scratch/runs/arcagi/exp010'),
                    working_dir=PosixPath('/opt/scratch/runs/arcagi/exp010/wandb'),
                    capture_console=True,
                    replay_startup_logs=False,
                    allow_startup_failure=True,
                    ingestion=WandbIngestion(
                        init_timeout_sec=30.0,
                        service_wait_sec=30.0,
                        flush_interval_sec=15.0,
                        system_metrics=True,
                        system_metrics_interval_sec=60.0
                    ),
                    metric_step_metrics={},
                    run_config={},
                    notes=''
                ),
                'metrics': FileTracker.Config(
                    base_dir=PosixPath('/opt/scratch/runs/arcagi/exp010'),
                    working_dir=PosixPath('/opt/scratch/runs/arcagi/exp010/metrics.json'),
                    capture_prefix='eval/'
                )},
            base_dir=PosixPath('/opt/scratch/runs/arcagi/exp010'),
            working_dir=PosixPath('/opt/scratch/runs/arcagi/exp010')
        ),
        max_steps=2_380_000,
        max_epochs=inf,
        max_time=inf,
        max_time_kind='wall',
        max_eval_time=inf,
        eval_stop_on_time_limit=False,
        num_steps_eval=10_000,
        num_steps_log=100,
        early_train_log_steps=100,
        eval_every_epoch=False,
        eval_extras_every_eval=False,
        eval_warmup_batches=0,
        eval_only=False,
        restore_rng_state=True,
        seed=0,
        model_torch_seed=torch.random.manual_seed,
        num_steps_garbage_collect=inf,
        mesh_dim_model_seed='pp',
        mesh_dim_data_seed=None,
        runtime=MultiProcess.Config(
            device=None,
            deterministic=False,
            float32_matmul_precision=None,
            backend=None,
            process_group_timeout_sec=None,
            mesh_topology={'dp': 8, 'pp': 1, 'tp': 1}
        ),
        global_batch_size=768
    )
