Add assert_not_both_not_none (fixes #393)

Author: Sorata-kanda

chex/_src/asserts.py

@@ -25,6 +25,7 @@
 
 from chex._src import asserts_internal as _ai
 from chex._src import pytypes
+from chex._src import asserts
 import jax
 from jax.experimental import checkify
 import jax.numpy as jnp
@@ -248,6 +249,20 @@ def assert_not_both_none(first: Any, second: Any) -> None:
     raise AssertionError(
         "At least one of the arguments must be different from `None`.")
 
+@_static_assertion
+def assert_not_both_not_none(first: Any, second: Any) -> None:
+    """Checks that not both arguments are non-None.
+
+    Args:
+      first: A first object.
+      second: A second object.
+
+    Raises:
+      AssertionError: If both ``first`` and ``second`` are not None.
+    """
+    if first is not None and second is not None:
+        raise AssertionError(
+            "At most one of the arguments may be different from `None`.")
 
 @_static_assertion
 def assert_exactly_one_is_none(first: Any, second: Any) -> None:

chex/_src/asserts_test.py

@@ -1860,7 +1860,15 @@ def test_assert_equal_fail(self, first, second):
     with self.assertRaises(AssertionError):
       asserts.assert_equal(first, second)
 
+class NotNoneAssertionsTest(parameterized.TestCase):
 
+    def test_assert_not_both_not_none(self):
+        asserts.assert_not_both_not_none(None, None)
+        asserts.assert_not_both_not_none(1, None)
+        asserts.assert_not_both_not_none(None, 1)
+
+        with self.assertRaises(AssertionError):
+            asserts.assert_not_both_not_none(1, 2)
 class IsDivisibleTest(parameterized.TestCase):
 
   def test_assert_is_divisible(self):