feat: make Rain hardware bias beta-independent

1 files changed, 20 insertions, 0 deletions

diff --git a/experiments/rain_ep_layer_adapter_smoke.py b/experiments/rain_ep_layer_adapter_smoke.py
index 8959217..24c9e0a 100644
--- a/experiments/rain_ep_layer_adapter_smoke.py
+++ b/experiments/rain_ep_layer_adapter_smoke.py
@@ -125,12 +125,32 @@ def main() -> None:
     assert innovation.debiaser.neutral_observations == 64
     assert constant.debiaser.neutral_observations == 64
     assert all(not value.requires_grad for value in innovation_used.values())
+
+    positive_field = RainLayerStateCorrector(
+        mode="raw", bias_ratio=2e-4, bias_normalization="first_state", seed=47)
+    negative_field = RainLayerStateCorrector(
+        mode="raw", bias_ratio=2e-4, bias_normalization="first_state", seed=47)
+    positive_field.apply(first, second, layer_names)
+    negative_second = {
+        name: value - clean_difference[name] for name, value in first.items()
+    }
+    negative_field.apply(first, negative_second, layer_names)
+    positive_bias = positive_field._measure(
+        [first[name] for name in layer_names])[1]
+    negative_bias = negative_field._measure(
+        [first[name] for name in layer_names])[1]
+    beta_independent_bias = all(
+        torch.equal(positive, negative)
+        for positive, negative in zip(positive_bias, negative_bias)
+    )
+    assert beta_independent_bias
     print({
         "oracle_parameter_gradient_relative_error": oracle_relative_error,
         "innovation_heldout_state_residual_rms": innovation_error,
         "constant_heldout_state_residual_rms": constant_error,
         "matched_neutral_observations": 64,
         "extra_equilibrium_phases": 0,
+        "first_state_bias_bitwise_independent_of_beta_sign": True,
         "requires_grad": False,
     })