ArcTrainLoop(
                study_name='arcagi2',
                experiment_name='exp000',
                base_dir='/opt/scratch',
                working_dir=PosixPath('/opt/scratch/runs/arcagi2/exp000'),
                doc='',
                step=ArcTrainStep.Config(
                │       model=ArcModelConfig(
                │       │       channels_in=512,
                │       │       num_layers=2,
                │       │       embedding=GridEmbedding.Config(
                │       │               channels_in=12,
                │       │               channels_out=512,
                │       │               grid_shape=(900,),
                │       │               channels=[],
                │       │               init_weight=<function priml.baselines.sudoku.embedding.corrected_scaled_normal at 0xdefacedeface>
                │       │       ),
                │       │       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,
                │       │       │               norm_qk=None,
                │       │       │               share_qk_norm=True,
                │       │       │               norm_out=None,
                │       │       │               split_qkv_projection=False,
                │       │       │               init_weight=<function priml.baselines.sudoku.model.corrected_fan_in_normal at 0xdefacedeface>,
                │       │       │               depth_index=(),
                │       │       │               attn_kernel=SdpaFused.Config()
                │       │       │       ),
                │       │       │       ffn=SwiGLU.Config(
                │       │       │               channels_in=512,
                │       │       │               channels_out=512,
                │       │       │               channels_hidden=1_536,
                │       │       │               expansion=2.6666666666666665,
                │       │       │               round_to=256,
                │       │       │               gate=True,
                │       │       │               bias=False,
                │       │       │               split_gate_projection=False,
                │       │       │               norm=None,
                │       │       │               act=<function priml.model.swiglu.silu at 0xdefacedeface>,
                │       │       │               depth_index=(),
                │       │       │               init_weight=<function priml.baselines.sudoku.model.corrected_fan_in_normal at 0xdefacedeface>,
                │       │       │               init_weight_out=<function priml.baselines.sudoku.model.corrected_fan_in_normal at 0xdefacedeface>,
                │       │       │               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=(),
                │       │       │       rope=RoPE.Config(
                │       │       │               channels_head=64,
                │       │       │               reduction_mode='cat',
                │       │       │               frequencies=HuggingFaceFrequencies.Config(
                │       │       │                       base=10000.0
                │       │       │               ),
                │       │       │               dtype=None
                │       │       │       )
                │       │       ),
                │       │       recurrence=DeepRecurrence.Config(
                │       │               slow_cycles=3,
                │       │               fast_cycles=4
                │       │       ),
                │       │       prefix=PuzzleEmbedding.Config(
                │       │               channels_in=-1,
                │       │               channels_out=512,
                │       │               num_puzzles=1_191_727,
                │       │               num_tokens=16,
                │       │               batch_size=256,
                │       │               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=<function priml.baselines.sudoku.model.corrected_fan_in_normal at 0xdefacedeface>,
                │       │               init_bias=<function torch.nn.init.zeros_ at 0xdefacedeface>,
                │       │               shard=None
                │       │       ),
                │       │       init_latent=<function priml.baselines.sudoku.model.corrected_unit_normal at 0xdefacedeface>,
                │       │       halt_outputs=2,
                │       │       halt_init_bias=-5.0,
                │       │       dtype=None,
                │       │       rope=None,
                │       │       rope_grid_shape=(),
                │       │       compile_core=None,
                │       │       dtype_latent_init=None,
                │       │       checkpoint_core=False,
                │       │       block_checkpoint_fraction=0.0
                │       ),
                │       optimizer=AdamATan2.Config(
                │               lr=0.0001,
                │               betas=(0.9, 0.95),
                │               weight_decay=0.1
                │       ),
                │       loss=SimpleLoss.Config(
                │               loss_fn=<function priml.loss.simple_loss.bce_with_logits at 0xdefacedeface>,
                │               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(<class 'functools.partial'>, <function priml.math.schedules.constant at 0xdefacedeface>),
                │       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=PartialConfig(<class 'functools.partial'>, <function torch.compile at 0xdefacedeface>, fullgraph=True),
                │       ema=NoEMA.Config(),
                │       gradient_clip_norm=inf,
                │       skip_step_on_nonfinite_grad=False,
                │       device_init='eager',
                │       dtype_autocast=torch.bfloat16,
                │       autocast_cache_enabled=False,
                │       total_train_steps=541_580,
                │       warmup_steps=2_000,
                │       lr_min_ratio=1.0,
                │       pool=AtomicPool.Config(
                │               batch_size=256,
                │               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,
                │               dtype=None
                │       ),
                │       label_smoothing=0.0,
                │       ignore_label_id=-100,
                │       use_ema=True,
                │       ema_decay=0.999,
                │       ema_warmup_steps=2_000,
                │       sparse_optimizer=SignSGD.Config(
                │               lr=0.01,
                │               weight_decay=0.1,
                │               aggregate_distributed=True
                │       )
                ),
                dataset=Arc2Data.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={1: 1.0},
                                        translation_prob=1.0,
                                        scale_prob=1.0
                                ),
                                spatial_eval_views=False,
                                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/arcagi2/arc2concept-aug-1000'),
                        batch_size=256,
                        eval_batch_size=None,
                        device='auto',
                        seed=0,
                        num_tasks=None,
                        num_eval_tasks=None,
                        epochs_per_iter=4
                ),
                metrics_train={},
                metrics_eval={       '': PassK.Config(
                                base_dir='/opt/scratch',
                                working_dir=PosixPath('/opt/scratch/datasets/arcagi2/arc2concept-aug-1000'),
                                pass_ks=(1, 2, 5, 10, 100, 1_000)
                        )},
                checkpointer=Checkpointer.Config(
                        base_dir=PosixPath('/opt/scratch/runs/arcagi2/exp000'),
                        working_dir=PosixPath('/opt/scratch/runs/arcagi2/exp000/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/arcagi2/exp000'),
                        working_dir=PosixPath('/opt/scratch/runs/arcagi2/exp000/profiling'),
                        heartbeat_interval_sec=20.0,
                        fault_dump_interval_sec=40.0,
                        cuda_events=False
                ),
                tracker=None,
                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=False,
                eval_warmup_batches=1,
                eval_only=False,
                restore_rng_state=True,
                seed=0,
                num_steps_garbage_collect=inf,
                mesh_dim_model_seed='pp',
                mesh_dim_data_seed='dp',
                runtime=MultiProcess.Config(
                        device='auto',
                        deterministic=False,
                        float32_matmul_precision=None,
                        backend=None,
                        mesh_topology={'dp': -1, 'pp': 1, 'tp': 1}
                )
        )
