MazeTrainLoop(
        study_name='maze',
        experiment_name='exp002',
        base_dir='/opt/scratch',
        working_dir=PosixPath('/opt/scratch/runs/maze/exp002'),
        doc='',
        step=BoundaryLossTrainStep.Config(
        │   model=GridSolver.Config(
        │   │   channels_in=512,
        │   │   num_layers=4,
        │   │   embedding=GridEmbedding.Config(
        │   │       channels_in=6,
        │   │       channels_out=512,
        │   │       grid_shape=(30, 30),
        │   │       channels=[],
        │   │       init_weight=functools.partial(priml.model.init.fan_in_truncated_normal, variance_correction=True),
        │   │       token_table=None,
        │   │       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=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=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=()
        │   │   ),
        │   │   recurrence=SingleLatentRecurrence.Config(
        │   │       slow_cycles=1,
        │   │       fast_cycles=8,
        │   │       inner_grad_loops=5
        │   │   ),
        │   │   prefix=SparsePuzzleEmbedding.Config(
        │   │       channels_in=-1,
        │   │       channels_out=512,
        │   │       num_puzzles=1,
        │   │       num_tokens=1,
        │   │       batch_size=128,
        │   │       init_std=0.0,
        │   │       embed_scale=-1.0,
        │   │       dtype=torch.bfloat16,
        │   │       dtype_scale=torch.float32
        │   │   ),
        │   │   num_prefix_tokens=1,
        │   │   vocab_size=6,
        │   │   head=Linear.Config(
        │   │       channels_in=512,
        │   │       channels_out=6,
        │   │       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=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=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=NoParallel.Config(device='cuda'),
        │   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=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.0001,
        │       weight_decay=1.0,
        │       aggregate_distributed=True
        │   ),
        │   optimizer_lr_policy=priml.baselines.arcagi.train_step.scale_body_optimizer_lr,
        │   pool=AtomicPool.Config(
        │       batch_size=128,
        │       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=901,
        │       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=128),
        │   signals=None,
        │   ignore_label_id=0,
        │   total_train_steps=24_000,
        │   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=False,
        │   log_body_norms=True,
        │   norm_log_interval=10,
        │   boundary=CanonicalBoundary.Config(weight=0.1, grid_shape=(30, 30))
        ),
        dataset=PuzzleDataset.Config(
            spec=SudokuSpec(grid_shape=(30, 30), box_shape=(3, 3), vocab_size=6),
            base_dir='/opt/scratch',
            working_dir=PosixPath('/opt/scratch/datasets/maze/hard-1k-canonical'),
            batch_size=128,
            eval_batch_size=None,
            device='cuda',
            augment=False,
            augment_digits_only=False,
            augment_seed=None,
            seed=0,
            num_instances=None,
            eval_num_instances=None,
            eval_instance_indices=(),
            max_samples=None
        ),
        metrics_train={},
        metrics_eval={},
        checkpointer=Checkpointer.Config(
            base_dir=PosixPath('/opt/scratch/runs/maze/exp002'),
            working_dir=PosixPath('/opt/scratch/runs/maze/exp002/checkpoints'),
            filename='step_{step:08d}.pt',
            save_every=500,
            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/maze/exp002'),
            working_dir=PosixPath('/opt/scratch/runs/maze/exp002/profiling'),
            heartbeat_interval_sec=20.0,
            fault_dump_interval_sec=40.0,
            cuda_events=False
        ),
        tracker=TrackerList.Config(
            trackers={   'wandb': WandbTracker.Config(
                    project='maze',
                    name='',
                    run_id='',
                    group='',
                    mode='online',
                    base_dir=PosixPath('/opt/scratch/runs/maze/exp002'),
                    working_dir=PosixPath('/opt/scratch/runs/maze/exp002/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/maze/exp002'),
                    working_dir=PosixPath('/opt/scratch/runs/maze/exp002/metrics.json'),
                    capture_prefix='eval/'
                )},
            base_dir=PosixPath('/opt/scratch/runs/maze/exp002'),
            working_dir=PosixPath('/opt/scratch/runs/maze/exp002')
        ),
        max_steps=24_000,
        max_epochs=inf,
        max_time=3600.0,
        max_time_kind='train',
        max_eval_time=3600.0,
        eval_stop_on_time_limit=False,
        num_steps_eval=500,
        num_steps_log=10,
        early_train_log_steps=100,
        eval_every_epoch=False,
        eval_extras_every_eval=False,
        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=SingleProcess.Config(
            device='cuda',
            deterministic=False,
            float32_matmul_precision=None
        )
    )
