diff --git a/tests/test_torch_preprocessing/test_torch_squashing_scaler.py b/tests/test_torch_preprocessing/test_torch_squashing_scaler.py index e63c499b3..c290ba8f0 100644 --- a/tests/test_torch_preprocessing/test_torch_squashing_scaler.py +++ b/tests/test_torch_preprocessing/test_torch_squashing_scaler.py @@ -439,3 +439,23 @@ def test__blocked_transform_does_not_write_to_its_input( _run_with_budgets(patch, _FORCE_BLOCKED, x, 64, 3.0) _assert_bit_identical("input tensor", x, before) + + +@pytest.mark.parametrize("outlier", [1000.0, -1000.0]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test__call__low_precision_outlier_does_not_overflow( + outlier: float, dtype: torch.dtype +) -> None: + # Both low-precision dtypes take the upcast path, so both need the cover. + scaler = TorchSquashingScaler() + x = torch.tensor( + [0, 1, 1, 1, 2, outlier], + dtype=dtype, + ).unsqueeze(-1) + + out = scaler(x, num_train_rows=5) + + assert out.dtype == dtype + assert torch.isfinite(out).all() + assert 2.5 < abs(out[-1, 0].item()) <= 3.0 + assert out[-1, 0].item() * outlier > 0