Arc2TrainLoop(
        study_name='arcagi',
        experiment_name='exp008',
        base_dir='/opt/scratch',
        working_dir=PosixPath('/opt/scratch/runs/arcagi/exp008'),
        doc='',
        step=PooledActTrainStep.Config(
        │   model=GridSolver.Config(
        │   │   channels_in=512,
        │   │   num_layers=4,
        │   │   embedding=GridEmbedding.Config(
        │   │       channels_in=12,
        │   │       channels_out=512,
        │   │       grid_shape=(900,),
        │   │       channels=[],
        │   │       init_weight=functools.partial(priml.model.init.fan_in_truncated_normal, variance_correction=True),
        │   │       token_table=None,
        │   │       embed_scale=-1.0
        │   │   ),
        │   │   block=RotaryBlock.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=None,
        │   │   │       proj_out=None,
        │   │   │       init_weight=functools.partial(priml.model.init.fan_in_truncated_normal, variance_correction=True),
        │   │   │       depth_index=(),
        │   │   │       attn_kernel=SdpaFused.Config()
        │   │   │   ),
        │   │   │   ffn=ConvSwiGLU.Config(
        │   │   │       channels_in=512,
        │   │   │       channels_out=512,
        │   │   │       channels_hidden=1_536,
        │   │   │       expansion=2.6666666666666665,
        │   │   │       round_to=256,
        │   │   │       kernel_size=2,
        │   │   │       causal=False,
        │   │   │       bias=False,
        │   │   │       short_conv=priml.baselines.arcagi.model.depthwise_shift,
        │   │   │       norm=RMSNorm.Config(
        │   │   │           channels_in=1_536,
        │   │   │           channels_out=1_536,
        │   │   │           eps=1e-06,
        │   │   │           elementwise_affine=False,
        │   │   │           device=None,
        │   │   │           dtype=None
        │   │   │       ),
        │   │   │       init_weight=functools.partial(priml.model.init.fan_in_truncated_normal, variance_correction=True)
        │   │   │   ),
        │   │   │   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=(),
        │   │   │   rope=RoPE.Config(
        │   │   │       channels_head=64,
        │   │   │       reduction_mode='cat',
        │   │   │       frequencies=HuggingFaceFrequencies.Config(base=10000.0),
        │   │   │       dtype=torch.float32
        │   │   │   )
        │   │   ),
        │   │   recurrence=SingleLatentRecurrence.Config(
        │   │       slow_cycles=2,
        │   │       fast_cycles=6,
        │   │       inner_grad_loops=0
        │   │   ),
        │   │   prefix=PuzzleEmbedding.Config(
        │   │       channels_in=-1,
        │   │       channels_out=512,
        │   │       num_puzzles=1_191_727,
        │   │       num_tokens=16,
        │   │       batch_size=96,
        │   │       init_std=0.0,
        │   │       embed_scale=-1.0,
        │   │       dtype=torch.bfloat16,
        │   │       dtype_scale=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=<built-in function linear>,
        │   │       shard=None
        │   │   ),
        │   │   init_latent=functools.partial(priml.model.init.truncated_normal, std=1.0, variance_correction=True),
        │   │   halt_outputs=2,
        │   │   halt_init_bias=-5.0,
        │   │   dtype=None,
        │   │   rope=None,
        │   │   rope_grid_shape=(),
        │   │   compile_core=CoreCompile.Config(unit='core', mode='default', fullgraph=True),
        │   │   dtype_latent_init=None,
        │   │   checkpoint_core=False,
        │   │   block_checkpoint_fraction=0.0
        │   ),
        │   optimizer=CompositeOptimizer.Config(
        │       optimizers=[
        │           PartialConfig(
        │               <class 'torch.optim.adamw.AdamW'>,
        │               lr=0.0001,
        │               betas=(0.9, 0.95),
        │               weight_decay=1.0
        │           ),
        │           Muon.Config(
        │               lr=0.005,
        │               weight_decay=0.02,
        │               momentum=0.6,
        │               nesterov=True,
        │               ns_coefficients=(3.4445, -4.775, 2.0315),
        │               eps=1e-07,
        │               ns_steps=3,
        │               reference_numerics=False,
        │               adjust_lr_fn=priml.optimizers.muon.adjust_lr_original,
        │               momentum_update_fn=None,
        │               matrix_signum_fn=None,
        │               ensemble_dims=0
        │           ),
        │           Muon.Config(
        │               lr=0.005,
        │               weight_decay=0.02,
        │               momentum=0.6,
        │               nesterov=True,
        │               ns_coefficients=(3.4445, -4.775, 2.0315),
        │               eps=1e-07,
        │               ns_steps=3,
        │               reference_numerics=False,
        │               adjust_lr_fn=priml.optimizers.muon.adjust_lr_original,
        │               momentum_update_fn=None,
        │               matrix_signum_fn=None,
        │               ensemble_dims=1
        │           )
        │       ],
        │       select=[
        │           complement(excluding(Muon.eligible_tensor, 'embed', 'head', 'register_tokens')),
        │           with_ndim(excluding(Muon.eligible_tensor, 'embed', 'head', 'register_tokens'), 2),
        │           with_ndim(excluding(Muon.eligible_tensor, 'embed', 'head', 'register_tokens'), 3)
        │       ],
        │       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=ArcDataParallel.Config(
        │       mesh_dim='dp',
        │       bucket_cap_mb=25,
        │       find_unused_parameters=False,
        │       gradient_as_bucket_view=True
        │   ),
        │   model_quantization=NoModelQuantization.Config(),
        │   activation_memoization=DefaultActivationStorage.Config(),
        │   compile=None,
        │   ema=EMA.Config(
        │       decay=0.999,
        │       update_after_step=2_000,
        │       update_every=1,
        │       track_buffers=False,
        │       shadow_kind='param_dict',
        │       warmup_seed=True,
        │       decay_schedule=priml.train.ema.constant_decay,
        │       select=priml.optimizers.parameter_filter.trainable,
        │       update_fn=None,
        │       buffer_update_fn=priml.train.ema.copy_buffers
        │   ),
        │   gradient_clip_norm=inf,
        │   skip_step_on_nonfinite_grad=False,
        │   device_init='eager',
        │   dtype_autocast=torch.bfloat16,
        │   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.scale_body_optimizer_lr,
        │   pool=AtomicPool.Config(
        │       batch_size=96,
        │       max_steps=16,
        │       halting=HaltTraining.Config(
        │           weight=0.5,
        │           exploration=SampledMinimum.Config(prob=0.1),
        │           seed=0
        │       ),
        │       feedback=None,
        │       start=LearnedStart.Config(),
        │       grid_len=900,
        │       seq_len=916,
        │       channels_hidden=512,
        │       steps_dtype=torch.int64,
        │       dtype=None
        │   ),
        │   pool_transition=AfterOptimizerPoolTransition.Config(),
        │   ema_initialization=None,
        │   gradient_reduction=None,
        │   metric_reporter=None,
        │   token_loss=StablemaxTokens.Config(
        │       log_normalizer=priml.math.loss.log_stablemax
        │   ),
        │   objective=None,
        │   reduction=MeanOverBatch.Config(batch_size=96),
        │   signals=EvalSignals.Config(per_step=True, color_offset=2),
        │   ignore_label_id=-100,
        │   total_train_steps=541_580,
        │   backward_loss_scale=1.0,
        │   lr_step_offset=0,
        │   warmup_steps=2_000,
        │   lr_min_ratio=1.0,
        │   grad_clip_foreach=True,
        │   emulate_precision_casts=True,
        │   log_body_norms=True,
        │   norm_log_interval=100
        ),
        dataset=Arc2PuzzleDataset.Config(
        │   augmentation=ArcAugmentation.Config(
        │       num_aug=1_000,
        │       seed=42,
        │       retries_factor=5,
        │       transform=ColorDihedral.Config(
        │           transforms=(0, 1, 2, 3, 4, 5, 6, 7),
        │           colors=(1, 2, 3, 4, 5, 6, 7, 8, 9),
        │           separator='|||'
        │       ),
        │       spec=ArcSpec(
        │           max_grid=30,
        │           vocab_pad=0,
        │           vocab_eos=1,
        │           vocab_color_offset=2,
        │           puzzle_id_separator='|||'
        │       ),
        │       spatial=SpatialAugmentation.Config(
        │           spec=ArcSpec(
        │               max_grid=30,
        │               vocab_pad=0,
        │               vocab_eos=1,
        │               vocab_color_offset=2,
        │               puzzle_id_separator='|||'
        │           ),
        │           train_scale_weights={2: 1.0},
        │           translation_prob=0.2,
        │           scale_prob=0.2
        │       ),
        │       spatial_eval_views=True,
        │       spatial_eval_scale=2
        │   ),
        │   spec=ArcSpec(
        │       max_grid=30,
        │       vocab_pad=0,
        │       vocab_eos=1,
        │       vocab_color_offset=2,
        │       puzzle_id_separator='|||'
        │   ),
        │   base_dir='/opt/scratch',
        │   working_dir=PosixPath('/opt/scratch/datasets/arc2concept-aug-1000-tr0p2-sc0p2-2w1p0-n1000-s42-spatialeval-v2'),
        │   num_puzzle_identifiers=1_191_727,
        │   batch_size=96,
        │   eval_batch_size=256,
        │   device=None,
        │   seed=0,
        │   rank=-1,
        │   num_replicas=-1,
        │   epochs_per_iter=4,
        │   max_samples=None,
        │   spatial_eval_views=2,
        │   source_dataset_dir=PosixPath('/opt/scratch/datasets/arc2concept-aug-1000-tr0p2-sc0p2-2w1p0-n1000-s42'),
        │   input_file_prefix=''
        ),
        metrics_train={},
        metrics_eval={   '': CanonicalPassK.Config(
        │       base_dir='/opt/scratch',
        │       working_dir=PosixPath('/opt/scratch/datasets/arc2concept-aug-1000-tr0p2-sc0p2-2w1p0-n1000-s42-spatialeval-v2'),
        │       pass_ks=(1, 2, 5, 10, 100, 1_000),
        │       spatial_views='non_spatial',
        │       max_views_per_input=0,
        │       per_step_acts=16,
        │       spec=ArcSpec(
        │           max_grid=30,
        │           vocab_pad=0,
        │           vocab_eos=1,
        │           vocab_color_offset=2,
        │           puzzle_id_separator='|||'
        │       ),
        │       transform=ColorDihedral.Config(
        │           transforms=(0, 1, 2, 3, 4, 5, 6, 7),
        │           colors=(1, 2, 3, 4, 5, 6, 7, 8, 9),
        │           separator='|||'
        │       ),
        │       rules=[
        │           StrictPass.Config(name='strict'),
        │           PerOutputPass.Config(name='per_output')
        │       ]
        │   ),
        │   'spatial_eq': CanonicalPassK.Config(
        │       base_dir='/opt/scratch',
        │       working_dir=PosixPath('/opt/scratch/datasets/arc2concept-aug-1000-tr0p2-sc0p2-2w1p0-n1000-s42-spatialeval-v2'),
        │       pass_ks=(1, 2, 5, 10, 100, 1_000),
        │       spatial_views='all',
        │       max_views_per_input=1_001,
        │       per_step_acts=16,
        │       spec=ArcSpec(
        │           max_grid=30,
        │           vocab_pad=0,
        │           vocab_eos=1,
        │           vocab_color_offset=2,
        │           puzzle_id_separator='|||'
        │       ),
        │       transform=ColorDihedral.Config(
        │           transforms=(0, 1, 2, 3, 4, 5, 6, 7),
        │           colors=(1, 2, 3, 4, 5, 6, 7, 8, 9),
        │           separator='|||'
        │       ),
        │       rules=[
        │           StrictPass.Config(name='strict'),
        │           PerOutputPass.Config(name='per_output')
        │       ]
        │   ),
        │   'spatial_big': CanonicalPassK.Config(
        │       base_dir='/opt/scratch',
        │       working_dir=PosixPath('/opt/scratch/datasets/arc2concept-aug-1000-tr0p2-sc0p2-2w1p0-n1000-s42-spatialeval-v2'),
        │       pass_ks=(1, 2, 5, 10, 100, 1_000),
        │       spatial_views='all',
        │       max_views_per_input=0,
        │       per_step_acts=16,
        │       spec=ArcSpec(
        │           max_grid=30,
        │           vocab_pad=0,
        │           vocab_eos=1,
        │           vocab_color_offset=2,
        │           puzzle_id_separator='|||'
        │       ),
        │       transform=ColorDihedral.Config(
        │           transforms=(0, 1, 2, 3, 4, 5, 6, 7),
        │           colors=(1, 2, 3, 4, 5, 6, 7, 8, 9),
        │           separator='|||'
        │       ),
        │       rules=[
        │           StrictPass.Config(name='strict'),
        │           PerOutputPass.Config(name='per_output')
        │       ]
            )},
        checkpointer=Checkpointer.Config(
            base_dir=PosixPath('/opt/scratch/runs/arcagi/exp008'),
            working_dir=PosixPath('/opt/scratch/runs/arcagi/exp008/checkpoints'),
            filename='step_{step:08d}.pt',
            save_every=5_000,
            keep_last_n=8,
            keep_every=40_000,
            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/exp008'),
            working_dir=PosixPath('/opt/scratch/runs/arcagi/exp008/profiling'),
            heartbeat_interval_sec=20.0,
            fault_dump_interval_sec=40.0,
            cuda_events=False
        ),
        tracker=TrackerList.Config(
            trackers={   'wandb': WandbTracker.Config(
                    project='trm',
                    name='',
                    run_id='',
                    group='',
                    mode='online',
                    base_dir=PosixPath('/opt/scratch/runs/arcagi/exp008'),
                    working_dir=PosixPath('/opt/scratch/runs/arcagi/exp008/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/exp008'),
                    working_dir=PosixPath('/opt/scratch/runs/arcagi/exp008/metrics.json'),
                    capture_prefix='eval/'
                ),
                'signals': SignalDumpTracker.Config(
                    base_dir=PosixPath('/opt/scratch/runs/arcagi/exp008'),
                    working_dir=PosixPath('/opt/scratch/runs/arcagi/exp008/signals_{global_step}.npz'),
                    spec=ArcSpec(
                        max_grid=30,
                        vocab_pad=0,
                        vocab_eos=1,
                        vocab_color_offset=2,
                        puzzle_id_separator='|||'
                    ),
                    keep_last_n=8,
                    keep_every=40_000
                )},
            base_dir=PosixPath('/opt/scratch/runs/arcagi/exp008'),
            working_dir=PosixPath('/opt/scratch/runs/arcagi/exp008')
        ),
        max_steps=541_580,
        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=True,
        eval_warmup_batches=1,
        eval_only=False,
        restore_rng_state=True,
        seed=0,
        model_torch_seed=None,
        num_steps_garbage_collect=inf,
        mesh_dim_model_seed='pp',
        mesh_dim_data_seed='dp',
        runtime=MultiProcess.Config(
            device=None,
            deterministic=False,
            float32_matmul_precision=None,
            backend=None,
            process_group_timeout_sec=None,
            mesh_topology={'dp': -1, 'pp': 1, 'tp': 1}
        )
    )
