From 761ca4d66d052f0f18fa49a3658cef41ceb4a65c Mon Sep 17 00:00:00 2001 From: David Kelley Date: Sat, 10 Oct 2026 15:49:47 -0700 Subject: [PATCH] Leave loss inputs intact when applying position weights PoissonMultinomialLoss scaled y_true and y_pred in place, so the trainer's train/valid r, r2 and spec (and the spec stop_stat) were computed on edge-downweighted tensors. Co-Authored-By: Claude Opus 5.5 --- src/baskerville/metrics.py | 4 ++-- src/tests/test_loss_weight.py | 9 +++++++++ 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/src/baskerville/metrics.py b/src/baskerville/metrics.py index c5ecb42..01bc8bd 100644 --- a/src/baskerville/metrics.py +++ b/src/baskerville/metrics.py @@ -332,8 +332,8 @@ def forward(self, y_pred, y_true, target_weights=None): self._setup_position_weights(seq_len, y_true.device) # apply position weights - y_true *= self.position_weights - y_pred *= self.position_weights + y_true = y_true * self.position_weights + y_pred = y_pred * self.position_weights # sum across lengths s_true = y_true.sum(dim=-1) # B x T diff --git a/src/tests/test_loss_weight.py b/src/tests/test_loss_weight.py index d4d6251..1cb5dc4 100644 --- a/src/tests/test_loss_weight.py +++ b/src/tests/test_loss_weight.py @@ -139,6 +139,15 @@ def test_scaling_all_weights_is_a_noop(loss_inputs): assert torch.allclose(one, ten) +def test_position_weights_leave_inputs_intact(loss_inputs): + """The trainer computes metrics from the same tensors after the loss.""" + y_pred, y_true = loss_inputs + loss_fn = PoissonMultinomialLoss(weight_range=10, weight_exp=8) + p, t = y_pred.clone(), y_true.clone() + loss_fn(p, t) + assert torch.equal(p, y_pred) and torch.equal(t, y_true) + + @pytest.mark.parametrize("label", ["poisson", "poisson_mn"]) def test_every_loss_takes_weights(loss_inputs, label): y_pred, y_true = loss_inputs