Fix EntropicPotentials's test (#343)

ott-jax · Mar 22, 2023 · 9bc5477 · 9bc5477
1 parent 5bbf4a1
commit 9bc5477
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/tests/problems/linear/potentials_test.py b/tests/problems/linear/potentials_test.py
@@ -177,8 +177,10 @@ def test_entropic_potentials_pnorm(
     )
 
     if p == 1.0:
-      with pytest.raises(AssertionError, match="Legendre transform not"):
-        z = potentials.transport(x_test, forward=forward)
+      # h_legendre not defined in this case, NaNs will be returned, see also
+      # https://github.com/ott-jax/ott/pull/340
+      z = potentials.transport(x_test, forward=forward)
+      np.testing.assert_array_equal(z, np.nan)
     else:
       if forward:
         z = potentials.transport(x_test, forward=forward)