SudokuTrainLoop(
        study_name='sudoku',
        experiment_name='exp001',
        base_dir='/opt/scratch',
        working_dir=PosixPath('/opt/scratch/runs/sudoku/exp001'),
        doc='',
        step=RoutedActTrainStep.Config(
        │   model=LatentRecursion.Config(
        │   │   vocab_size=11,
        │   │   puzzle_grid_shape=(81,),
        │   │   channels_in=512,
        │   │   num_layers=2,
        │   │   num_heads=8,
        │   │   slow_cycles=3,
        │   │   fast_cycles=4,
        │   │   q_head_outputs=2,
        │   │   num_puzzle_identifiers=1,
        │   │   puzzle_emb_len=16,
        │   │   puzzle_emb_ndim=512,
        │   │   puzzle_emb_batch_size=384,
        │   │   puzzle_emb_init_std=0.0,
        │   │   pos2d_grid_shape=(9, 9),
        │   │   pos2d_box_shape=(3, 3),
        │   │   pos2d_init_std=1.0,
        │   │   feedback_init_std=0.0,
        │   │   compile=True,
        │   │   dtype=torch.bfloat16,
        │   │   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=priml.baselines.sudoku.latent_recursion.truncated_normal_corrected,
        │   │   │       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=RMSNorm.Config(
        │   │   │           channels_in=1_536,
        │   │   │           channels_out=1_536,
        │   │   │           eps=1e-06,
        │   │   │           elementwise_affine=False,
        │   │   │           device=None,
        │   │   │           dtype=None
        │   │   │       ),
        │   │   │       act=priml.model.swiglu.silu,
        │   │   │       depth_index=(),
        │   │   │       init_weight=priml.baselines.sudoku.latent_recursion.truncated_normal_corrected,
        │   │   │       init_weight_out=priml.baselines.sudoku.latent_recursion.truncated_normal_corrected,
        │   │   │       up_proj=None,
        │   │   │       middle=None,
        │   │   │       down_proj=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=()
        │   │   )
        │   ),
        │   optimizer=CompositeOptimizer.Config(
        │       optimizers=[
        │           PartialConfig(
        │               <class 'torch.optim.adamw.AdamW'>,
        │               lr=0.0001,
        │               betas=(0.9, 0.95),
        │               weight_decay=1.0,
        │               fused=False
        │           ),
        │           Muon.Config(
        │               lr=0.02,
        │               weight_decay=0.005,
        │               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.02,
        │               weight_decay=0.005,
        │               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')),
        │           with_ndim(excluding(Muon.eligible_tensor, 'embed', 'head'), 2),
        │           with_ndim(excluding(Muon.eligible_tensor, 'embed', 'head'), 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=None),
        │   model_quantization=NoModelQuantization.Config(),
        │   activation_memoization=DefaultActivationStorage.Config(),
        │   compile=None,
        │   ema=EMA.Config(
        │       decay=0.9,
        │       update_after_step=5_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=1.0,
        │   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
        │   ),
        │   pool=AtomicPool.Config(
        │       batch_size=384,
        │       max_steps=16,
        │       halting=HaltTraining.Config(
        │           weight=0.05,
        │           exploration=SampledMinimum.Config(prob=0.1),
        │           seed=0
        │       ),
        │       feedback=None,
        │       start=LearnedStart.Config(),
        │       grid_len=81,
        │       seq_len=97,
        │       channels_hidden=512,
        │       steps_dtype=torch.int64,
        │       dtype=None
        │   ),
        │   routing=SingleChain.Config(),
        │   token_loss=PartialConfig(
        │       priml.baselines.sudoku.losses.token_cross_entropy,
        │       label_smoothing=0.2
        │   ),
        │   cardinality=None,
        │   release=None,
        │   ignore_label_id=-100,
        │   total_train_steps=62_000,
        │   warmup_steps=0,
        │   lr_min_ratio=0.0,
        │   grad_clip_foreach=True,
        │   emulate_precision_casts=True,
        │   log_body_norms=True,
        │   norm_log_interval=100
        ),
        dataset=PuzzleDataset.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,
            augment_digits_only=False,
            augment_seed=None,
            seed=0,
            num_instances=None,
            eval_num_instances=2_000,
            eval_instance_indices=(),
            max_samples=None
        ),
        metrics_train={},
        metrics_eval={'grid': GridAccuracy.Config(ignore_label_id=-100)},
        checkpointer=Checkpointer.Config(
            base_dir=PosixPath('/opt/scratch/runs/sudoku/exp001'),
            working_dir=PosixPath('/opt/scratch/runs/sudoku/exp001/checkpoints'),
            filename='step_{step:08d}.pt',
            save_every=1_000,
            keep_last_n=8,
            keep_every=10_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/sudoku/exp001'),
            working_dir=PosixPath('/opt/scratch/runs/sudoku/exp001/profiling'),
            heartbeat_interval_sec=20.0,
            fault_dump_interval_sec=40.0,
            cuda_events=False
        ),
        tracker=FileTracker.Config(
            base_dir=PosixPath('/opt/scratch/runs/sudoku/exp001'),
            working_dir=PosixPath('/opt/scratch/runs/sudoku/exp001/metrics.json'),
            capture_prefix='eval/'
        ),
        max_steps=62_000,
        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=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
        )
    )
