SudokuTrainLoop(
                study_name='sudoku',
                experiment_name='exp000',
                base_dir='/opt/scratch',
                working_dir=PosixPath('/opt/scratch/runs/sudoku/exp000'),
                doc='',
                step=SudokuTrainStep.Config(
                │       model=SudokuNet.Config(
                │       │       channels_in=512,
                │       │       num_layers=2,
                │       │       embedding=GridEmbedding.Config(
                │       │               channels_in=11,
                │       │               channels_out=512,
                │       │               grid_shape=(81,),
                │       │               channels=[
                │       │                       FactoredPositions.Config(
                │       │                               grid_shape=(9, 9),
                │       │                               box_shape=(3, 3),
                │       │                               channels_out=512,
                │       │                               init_std=1.0,
                │       │                               embed_scale=-1.0
                │       │                       )
                │       │               ],
                │       │               init_weight=functools.partial(<function priml.model.init.fan_in_truncated_normal at 0xdefacedeface>, variance_correction=True)
                │       │       ),
                │       │       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,
                │       │       │               norm_qk=None,
                │       │       │               share_qk_norm=True,
                │       │       │               norm_out=None,
                │       │       │               split_qkv_projection=False,
                │       │       │               init_weight=<function priml.model.init.kaiming_uniform 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.model.init.kaiming_uniform at 0xdefacedeface>,
                │       │       │               init_weight_out=<function priml.model.init.kaiming_uniform at 0xdefacedeface>,
                │       │       │               shard='colwise'
                │       │       │       ),
                │       │       │       norm1=RMSNorm.Config(
                │       │       │               channels_in=512,
                │       │       │               channels_out=512,
                │       │       │               eps=1e-06,
                │       │       │               elementwise_affine=False,
                │       │       │               device=None,
                │       │       │               dtype=None
                │       │       │       ),
                │       │       │       norm2=RMSNorm.Config(
                │       │       │               channels_in=512,
                │       │       │               channels_out=512,
                │       │       │               eps=1e-06,
                │       │       │               elementwise_affine=False,
                │       │       │               device=None,
                │       │       │               dtype=None
                │       │       │       ),
                │       │       │       prenorm=False,
                │       │       │       checkpoint=False,
                │       │       │       depth_index=()
                │       │       ),
                │       │       recurrence=None,
                │       │       prefix=None,
                │       │       num_prefix_tokens=0,
                │       │       vocab_size=11,
                │       │       head=Linear.Config(
                │       │               channels_in=512,
                │       │               channels_out=11,
                │       │               bias=False,
                │       │               depth_index=(),
                │       │               device=None,
                │       │               dtype=None,
                │       │               init_weight=functools.partial(<function priml.model.init.fan_in_truncated_normal at 0xdefacedeface>, variance_correction=True),
                │       │               init_bias=<function torch.nn.init.zeros_ at 0xdefacedeface>,
                │       │               shard=None
                │       │       ),
                │       │       init_latent=functools.partial(<function priml.model.init.truncated_normal at 0xdefacedeface>, std=1.0, variance_correction=True),
                │       │       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=CompositeOptimizer.Config(
                │               optimizers=[
                │                       PartialConfig(<class 'functools.partial'>, <class 'torch.optim.adamw.AdamW'>, lr=0.0001, betas=(0.9, 0.95), weight_decay=1.0),
                │                       Muon.Config(
                │                               lr=0.02,
                │                               weight_decay=0.1,
                │                               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=<function priml.optimizers.muon.adjust_lr_original at 0xdefacedeface>,
                │                               ensemble_dims=0
                │                       )
                │               ],
                │               select=[
                │                       complement(excluding(Muon.eligible_tensor, 'embed', 'head')),
                │                       excluding(Muon.eligible_tensor, 'embed', 'head')
                │               ],
                │               require_total=True,
                │               drop_empty=False
                │       ),
                │       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=NoParallel.Config(device=None),
                │       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=19_500,
                │       warmup_steps=0,
                │       lr_min_ratio=0.0,
                │       pool=None,
                │       label_smoothing=0.0,
                │       ignore_label_id=-100,
                │       use_ema=True,
                │       ema_decay=0.999,
                │       ema_warmup_steps=0
                ),
                dataset=SudokuData.Config(
                        spec=SudokuSpec(
                                grid_shape=(9, 9),
                                box_shape=(3, 3),
                                vocab_size=11
                        ),
                        base_dir='/opt/scratch',
                        working_dir=PosixPath('/opt/scratch/datasets/sudoku-extreme'),
                        batch_size=384,
                        eval_batch_size=None,
                        device=None,
                        augment=True,
                        seed=0,
                        augment_seed=None,
                        num_train_puzzles=None,
                        num_eval_puzzles=2_000
                ),
                metrics_train={},
                metrics_eval={'accuracy': GridAccuracy.Config(ignore_label_id=-100)},
                checkpointer=Checkpointer.Config(
                        base_dir=PosixPath('/opt/scratch/runs/sudoku/exp000'),
                        working_dir=PosixPath('/opt/scratch/runs/sudoku/exp000/checkpoints'),
                        filename='step_{step:08d}.pt',
                        save_every=1_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/sudoku/exp000'),
                        working_dir=PosixPath('/opt/scratch/runs/sudoku/exp000/profiling'),
                        heartbeat_interval_sec=20.0,
                        fault_dump_interval_sec=40.0,
                        cuda_events=False
                ),
                tracker=None,
                max_steps=19_500,
                max_epochs=inf,
                max_time=inf,
                max_time_kind='wall',
                max_eval_time=inf,
                eval_stop_on_time_limit=False,
                num_steps_eval=1_000,
                num_steps_log=10,
                early_train_log_steps=100,
                eval_every_epoch=True,
                eval_extras_every_eval=False,
                eval_warmup_batches=0,
                eval_only=False,
                restore_rng_state=True,
                seed=None,
                num_steps_garbage_collect=inf,
                mesh_dim_model_seed='pp',
                mesh_dim_data_seed='dp',
                runtime=SingleProcess.Config(
                        device=None,
                        deterministic=False,
                        float32_matmul_precision=None
                )
        )
