Add quantile_uni_extrapolate preprocessor preset

[LeoGrin]

Summary

Adds the quantile_uni_extrapolate preprocessor preset: instead of clamping quantile-transformed values at the [0, 1] boundary, it linearly extrapolates past the boundary for inputs outside the training range, better preserving out-of-distribution information.

Wired through both the CPU AdaptiveQuantileTransformer and the GPU TorchQuantileTransformer paths so the preset is GPU-eligible. CPU/GPU consistency tests parametrise the new preset across f16/f32/f64.

Motivation

v3 OOD-preprocessing checkpoints use quantile_uni_extrapolate in the regressor preprocessing recipe. Without this preset in the public package those checkpoints fail to load (pydantic ValidationError on the PreprocessorConfig.name literal). This unblocks loading them with the public release.

Changes (7 source + 3 test files, mirrors private #599)

• preprocessing/configs.py — register quantile_uni_extrapolate in the name literal
• preprocessing/steps/adaptive_quantile_transformer.py — extrapolate_ratio logic (CPU)
• preprocessing/steps/reshape_feature_distribution_step.py — preset wiring
• preprocessing/torch/{factory,steps,gpu_preprocessing_metadata,torch_quantile_transformer}.py — GPU path
• tests for CPU transform, torch transform, and CPU/GPU consistency

Test plan

• test_adaptive_quantile_transformer.py — 5/5 pass (incl. NaN-column & constant-feature edge cases)
• test_torch_quantile_transformer.py extrapolation suite — 3/3 pass (incl. matches_cpu_on_out_of_range_inputs)

Note: repo policy is "open an issue first" — happy to file/link one; opening this so the diff is reviewable.

🤖 Generated with Claude Code

---

Note

Medium Risk
Touches core preprocessing behavior for a new quantile preset and changes both CPU (sklearn) and GPU (torch) quantile paths, which can affect model inputs and CPU/GPU parity if edge cases slip through.

Overview
Adds a new preprocessor preset, quantile_uni_extrapolate, that preserves out-of-distribution signal by linearly extrapolating quantile-transformed values beyond the usual [0, 1] clamp (with configurable extrapolate_ratio, defaulted to 1.0 for the preset).

Wires this behavior through both the CPU AdaptiveQuantileTransformer (stores per-feature train min/max and extrapolates at transform time, skipping constant features) and the GPU pipeline (TorchQuantileTransformer + Torch*QuantileTransformerStep/factory), and marks the preset as GPU-eligible. Updates tests to cover extrapolation semantics, validation guards, and CPU/GPU consistency for the new preset.

^{Reviewed by Cursor Bugbot for commit 1569371. Bugbot is set up for automated code reviews on this repo. Configure here.}

[gemini-code-assist]

Code Review

This pull request adds a new quantile_uni_extrapolate preset to the preprocessing pipeline, enabling linear extrapolation for values outside the training range in both AdaptiveQuantileTransformer (CPU) and TorchQuantileTransformer (GPU). This is achieved by adding an extrapolate_ratio parameter and corresponding logic to the transform methods. Reviewers suggested adding input validation to ensure the extrapolate_ratio is non-negative and only used with uniform output distributions to avoid incorrect results.

[gemini-code-assist]

It is recommended to validate that extrapolate_ratio is non-negative and that it is only used when output_distribution is set to "uniform". Linear extrapolation as implemented here is not mathematically appropriate for a normal output distribution and would lead to incorrect results if accidentally configured that way.

        self.extrapolate_ratio = extrapolate_ratio
        if extrapolate_ratio is not None:
            if extrapolate_ratio < 0:
                raise ValueError("extrapolate_ratio must be non-negative.")
            if kwargs.get("output_distribution", "uniform") != "uniform":
                raise ValueError("extrapolate_ratio is only supported for output_distribution='uniform'.")

[gemini-code-assist]

Consider adding a validation check to ensure extrapolate_ratio is non-negative. While the current presets use 1.0, explicit validation prevents potential issues if the class is used with custom configurations in the future.

Suggested change

	self.extrapolate_ratio = extrapolate_ratio
	self.extrapolate_ratio = extrapolate_ratio
	if extrapolate_ratio is not None and extrapolate_ratio < 0:
	raise ValueError("extrapolate_ratio must be non-negative.")

[LeoGrin]

Thanks @gemini-code-assist — both addressed in b54cb09:

1. extrapolate_ratio non-negative + uniform-only (CPU) — added in AdaptiveQuantileTransformer.__init__: raises ValueError for negative ratio, and for output_distribution != "uniform" (linear extrapolation past [0,1] is only meaningful for the uniform output mapping, matching the preset's documented intent).
2. extrapolate_ratio non-negative (GPU) — added the same non-negative guard in TorchQuantileTransformer.__init__.

Added tests for both guard paths (test__extrapolate_ratio__validation_guards, test__torch_extrapolate_ratio__rejects_negative).

Also in this push, unrelated to the bot comments but worth flagging for review:

• quantile_uni_extrapolate now uses the n_samples // 5 tier (same as quantile_uni), not the coarse n//10 tier it had in the internal PR. Rationale: the preset is meant to be "the default quantile transform plus boundary extrapolation" — it should differ from quantile_uni by extrapolation only, not also by a coarser quantile grid. This is an intentional deviation from internal Update the device doc string to include multi-gpu inference. #599.
• GPU _apply_extrapolation norm consolidation: the internal port computed norm_below and norm_above as two separate full-array tensors; (x - x_max)/range + 1 is algebraically identical to (x - x_min)/range, so it now computes one tensor and reuses it (the CPU side was already consolidated this way in Update the device doc string to include multi-gpu inference. #599). Output-identical — CPU/GPU consistency tests unchanged.

[gemini-code-assist]

Thanks for the update, @LeoGrin. The added validation guards and the consolidation of the GPU _apply_extrapolation logic look correct and improve the robustness of the implementation. The decision to use the n_samples // 5 tier for quantile_uni_extrapolate to maintain consistency with the base quantile_uni preset is also well-reasoned. The updated tests appropriately cover the new validation logic and verify the GPU/CPU consistency. This looks ready for final review.

[FBPrior]

LGTM thanks!

[FBPrior]

Do we ever want to set this to sth other than one? I think it's useful to keep this parameter around because it's a pretty important parameter.