Assert for either is populated but not both

[alembcke]

I am trying to convert these lines of code to JAX/Chex format. So I essentially want to invert chex.assert_not_both_none to get something like chex.assert_not_both_not_none meaning that one or the other is populated, but not both. Is there already a way to do that using Chex, so is there already a way to invert existing asserts? Or could this one be added?

Editing to add, what I am looking for is an assert where neither being populated is ok, one being populated is ok, but both being populated raises an assert. Also, if this should be added, then I would be happy to give it a try as my first contribution to Chex.