diff options
Diffstat (limited to 'sdil')
| -rw-r--r-- | sdil/rain_ep_adapter.py | 20 |
1 files changed, 20 insertions, 0 deletions
diff --git a/sdil/rain_ep_adapter.py b/sdil/rain_ep_adapter.py index 81442aa..9c7ab8f 100644 --- a/sdil/rain_ep_adapter.py +++ b/sdil/rain_ep_adapter.py @@ -109,6 +109,20 @@ class RainGradientCorrector: ) @torch.no_grad() + def observe_neutral(self, local_states: Iterable[Tensor]) -> None: + """Fit the local bias field from one instruction-off observation.""" + if self.mode not in {"constant", "innovation"}: + raise ValueError( + "neutral predictor observations require constant or innovation mode") + local_states = list(local_states) + if any(value.requires_grad for value in local_states): + raise ValueError("Rain adapter received a requires-grad tensor") + bases, bias = self.bias.measure(local_states) + if self.debiaser is None: + self._initialize_debiaser(local_states) + self.debiaser.update_neutral(bases, bias, self.predictor_rate) + + @torch.no_grad() def apply(self, clean: Iterable[Tensor], local_states: Iterable[Tensor]) -> list[Tensor]: clean = list(clean) local_states = list(local_states) @@ -187,3 +201,9 @@ def attach_to_rain_estimator(estimator, corrector: RainGradientCorrector): estimator.sdil_corrector = corrector return estimator + +@torch.no_grad() +def observe_rain_neutral(estimator, corrector: RainGradientCorrector) -> None: + """Expose one label-free Rain equilibrium to the local predictor.""" + local_states = [updater.grad() for updater in estimator._param_updaters] + corrector.observe_neutral(local_states) |