Bump ruff line-length to 120 and reformat

Rejoins lines that only wrapped because they exceeded the old 88-char
limit; ruff check and the full test suite (725 passed) are unaffected.

tests/test_train.py

@@ -74,9 +74,7 @@ def test_wandb_run_config_includes_full_cfg_and_param_counts():
         "stage1_model": {"generator": "flow"},
         "stage2_model": {"generator": "wgan"},
     }
-    wcfg = _wandb_run_config(
-        cfg, model_config={"pdg_vocab": 3}, param_counts={"stage1": 100}
-    )
+    wcfg = _wandb_run_config(cfg, model_config={"pdg_vocab": 3}, param_counts={"stage1": 100})
     assert wcfg["train"] == {"lr": 3e-4}
     assert wcfg["stage1_model"] == {"generator": "flow"}
     assert wcfg["stage2_model"] == {"generator": "wgan"}
@@ -162,9 +160,7 @@ def test_type_repr_shapes_and_values(target):
     expected_width = PARTICLE_PHYS_DIM if target == "physical" else emb_dim
     assert repr_.shape == (B, K, expected_width)
     if target == "physical":
-        assert torch.equal(
-            repr_, sec_cont[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM]
-        )
+        assert torch.equal(repr_, sec_cont[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM])
     if target == "onehot":
         assert torch.all(repr_.sum(-1) == 1.0)
 
@@ -180,9 +176,7 @@ def test_type_repr_shapes_and_values(target):
         ("embedding", "wgan"),
     ],
 )
-def test_assemble_stage2_ar_target_matches_assemble_stage2_real_flattened(
-    target, generator
-):
+def test_assemble_stage2_ar_target_matches_assemble_stage2_real_flattened(target, generator):
     """Regression test tying the refactor together: _assemble_stage2_real is
     now defined as _assemble_stage2_ar_target(...).flatten(1)."""
     B, emb_dim = 4, 6
@@ -192,12 +186,8 @@ def test_assemble_stage2_ar_target_matches_assemble_stage2_real_flattened(
     if target == "embedding":
         cond_enc.pdg_emb = torch.nn.Embedding(emb_dim, emb_dim)
     particle_type_cfg = {"target": target}
-    flat = _assemble_stage2_real(
-        sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim
-    )
-    unflat = _assemble_stage2_ar_target(
-        sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim
-    )
+    flat = _assemble_stage2_real(sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim)
+    unflat = _assemble_stage2_ar_target(sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim)
     assert torch.equal(unflat.flatten(1), flat)
 
 
@@ -206,9 +196,7 @@ def test_assemble_stage2_ar_inputs_shapes_and_history_feat_width():
     sec_cont = torch.randn(B, K_MAX, SEC_SLOT_DIM)
     sec_type_idx = torch.randint(0, emb_dim, (B, K_MAX))
     cond_enc = torch.nn.Module()
-    out = _assemble_stage2_ar_inputs(
-        sec_cont, sec_type_idx, {"target": "physical"}, cond_enc, emb_dim
-    )
+    out = _assemble_stage2_ar_inputs(sec_cont, sec_type_idx, {"target": "physical"}, cond_enc, emb_dim)
     assert out["history_feat"].shape == (B, K_MAX, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
     assert out["has_prev"].shape == (B, K_MAX)
     assert out["remaining_frac"].shape == (B, K_MAX)
@@ -221,9 +209,7 @@ def test_relax_onehot_type_slice_grad_probe_populates_both_norms():
     B, k_max, cont_dim, type_dim = 4, K_MAX, CONT_SLOT_DIM, 6
     x_flat = torch.randn(B, k_max * (cont_dim + type_dim), requires_grad=True)
     grad_probe: dict[str, float] = {}
-    out = _relax_onehot_type_slice(
-        x_flat, k_max, cont_dim, type_dim, tau=0.5, grad_probe=grad_probe
-    )
+    out = _relax_onehot_type_slice(x_flat, k_max, cont_dim, type_dim, tau=0.5, grad_probe=grad_probe)
     out.sum().backward()
     assert grad_probe["cont"] >= 0.0
     assert grad_probe["type"] >= 0.0
@@ -324,9 +310,7 @@ def _fake_batches(n_batches, batch_size, seed=0):
         sec_cont = torch.randn(batch_size, K_MAX, SEC_SLOT_DIM, generator=g)
         proc_idx = torch.zeros(batch_size, dtype=torch.long)
         sec_type_idx = torch.zeros(batch_size, K_MAX, dtype=torch.long)
-        batches.append(
-            (cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx)
-        )
+        batches.append((cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx))
     return batches
 
 
@@ -409,17 +393,13 @@ def _run_train(cfg, out_dir, resume_path=None):
         ),
         (
             "stage2_onehot_target_wgan",
-            lambda cfg: cfg["stage2_model"].__setitem__(
-                "particle_type", {"target": "onehot", "lambda": 1.0}
-            ),
+            lambda cfg: cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
         ),
         (
             "stage2_onehot_target_flow",
             lambda cfg: (
                 cfg["stage2_model"].__setitem__("generator", "flow"),
-                cfg["stage2_model"].__setitem__(
-                    "particle_type", {"target": "onehot", "lambda": 1.0}
-                ),
+                cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
             ),
         ),
         (
@@ -427,9 +407,7 @@ def _run_train(cfg, out_dir, resume_path=None):
             lambda cfg: (
                 cfg["conditioning"]["particle"].__setitem__("type", "embedding"),
                 cfg["conditioning"]["material"].__setitem__("type", "embedding"),
-                cfg["stage2_model"].__setitem__(
-                    "particle_type", {"target": "embedding", "lambda": 1.0}
-                ),
+                cfg["stage2_model"].__setitem__("particle_type", {"target": "embedding", "lambda": 1.0}),
             ),
         ),
         (
@@ -438,18 +416,14 @@ def _run_train(cfg, out_dir, resume_path=None):
                 cfg["conditioning"]["particle"].__setitem__("type", "embedding"),
                 cfg["conditioning"]["material"].__setitem__("type", "embedding"),
                 cfg["stage2_model"].__setitem__("generator", "flow"),
-                cfg["stage2_model"].__setitem__(
-                    "particle_type", {"target": "embedding", "lambda": 1.0}
-                ),
+                cfg["stage2_model"].__setitem__("particle_type", {"target": "embedding", "lambda": 1.0}),
             ),
         ),
         (
             "ar_wgan_onehot",
             lambda cfg: (
                 cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
-                cfg["stage2_model"].__setitem__(
-                    "particle_type", {"target": "onehot", "lambda": 1.0}
-                ),
+                cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
             ),
         ),
         (
@@ -461,9 +435,7 @@ def _run_train(cfg, out_dir, resume_path=None):
             lambda cfg: (
                 cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
                 cfg["stage2_model"].__setitem__("generator", "flow"),
-                cfg["stage2_model"].__setitem__(
-                    "particle_type", {"target": "onehot", "lambda": 1.0}
-                ),
+                cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
             ),
         ),
         (
@@ -473,9 +445,7 @@ def _run_train(cfg, out_dir, resume_path=None):
                 cfg["conditioning"]["material"].__setitem__("type", "embedding"),
                 cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
                 cfg["stage2_model"].__setitem__("generator", "flow"),
-                cfg["stage2_model"].__setitem__(
-                    "particle_type", {"target": "embedding", "lambda": 1.0}
-                ),
+                cfg["stage2_model"].__setitem__("particle_type", {"target": "embedding", "lambda": 1.0}),
             ),
         ),
         (
@@ -491,9 +461,7 @@ def _run_train(cfg, out_dir, resume_path=None):
                 cfg["stage1_model"].__setitem__("generator", "wgan"),
                 cfg["stage2_model"].__setitem__("generator", "flow"),
                 cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
-                cfg["stage2_model"].__setitem__(
-                    "particle_type", {"target": "onehot", "lambda": 1.0}
-                ),
+                cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
             ),
         ),
     ],
@@ -585,9 +553,7 @@ def test_wgan_stage_trainer_skips_generator_step_when_no_grad_this_batch():
         ema_decay=0.0,
         steps_per_epoch=4,
     )
-    trainer = WGANStageTrainer(
-        spec, models["stage1"], critics["stage1"], torch.device("cpu")
-    )
+    trainer = WGANStageTrainer(spec, models["stage1"], critics["stage1"], torch.device("cpu"))
     assert trainer.model.n_sec_head is None
     batch = _fake_batches(1, 8)[0]
     stats = trainer.step(batch, torch.device("cpu"), global_step=1)
@@ -595,9 +561,7 @@ def test_wgan_stage_trainer_skips_generator_step_when_no_grad_this_batch():
 
 
 def test_flow_stage_trainer_ddpm_not_implemented_for_stage2():
-    spec = StageSpec(
-        name="stage2", is_stage2=True, generator="ddpm", ddpm_n_steps=50, ema_decay=0.0
-    )
+    spec = StageSpec(name="stage2", is_stage2=True, generator="ddpm", ddpm_n_steps=50, ema_decay=0.0)
     with pytest.raises(NotImplementedError):
         FlowDDPMStageTrainer(spec, torch.nn.Linear(1, 1), torch.device("cpu"))
 
@@ -608,9 +572,7 @@ def test_flow_stage_trainer_ddpm_not_implemented_for_stage2():
 @pytest.mark.parametrize("teacher_forcing", ["always", "scheduled", "never"])
 @pytest.mark.parametrize("history", ["markov", "attention"])
 @pytest.mark.parametrize("stage2_generator", ["wgan", "flow"])
-def test_build_stage_trainers_ar_scheduled_and_attention_step_runs(
-    teacher_forcing, history, stage2_generator
-):
+def test_build_stage_trainers_ar_scheduled_and_attention_step_runs(teacher_forcing, history, stage2_generator):
     """v0.3.0 step 7: history='attention' and teacher_forcing in
     {'scheduled', 'never'} must actually train — a stage-2 AR trainer.step()
     must run and produce a finite loss, for every {history} x
@@ -629,9 +591,7 @@ def test_build_stage_trainers_ar_scheduled_and_attention_step_runs(
     model_config = _model_config(cfg)
     models = build_models(model_config)
     critics = build_critics(model_config)
-    trainers = build_stage_trainers(
-        cfg, models, critics, torch.device("cpu"), total_train_batches=4
-    )
+    trainers = build_stage_trainers(cfg, models, critics, torch.device("cpu"), total_train_batches=4)
     trainer = trainers["stage2"]
     batch = _fake_batches(1, 4)[0]
     stats = trainer.step(batch, torch.device("cpu"), global_step=1)
@@ -665,9 +625,7 @@ def test_train_end_to_end_ar_attention_history_scheduled_teacher_forcing(
         with open(out_dir / "metrics.csv", newline="") as f:
             rows = list(csv.DictReader(f))
         assert len(rows) == cfg["train"]["epochs"]
-        loss_col = (
-            "stage2/train/g_loss" if stage2_generator == "wgan" else "stage2/train/loss"
-        )
+        loss_col = "stage2/train/g_loss" if stage2_generator == "wgan" else "stage2/train/loss"
         assert all(math.isfinite(float(r[loss_col])) for r in rows)