Allow nan equality in unit test

google-deepmind · Jan 14, 2025 · 466ee20 · 466ee20
1 parent b832bb4
commit 466ee20
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/torax/tests/test_run_simulation_main.py b/torax/tests/test_run_simulation_main.py
@@ -237,6 +237,8 @@ def check_equality(ds1: xr.Dataset, ds2: xr.Dataset):
               # for other tests.
               # rtol=0.0,
               atol=5.0e-5,
+              # This is required to allow one of psi_right_grad_constraint and psi_right_constraint to be None
+              equal_nan=True,
           ):
             diff = ov - gv
             max_diff = np.abs(diff).max()