Method where

ot/backend.py:1771–1775

(self, condition, x=None, y=None)

        return a

    def where(self, condition, x=None, y=None):
        if x is None and y is None:
            return jnp.where(condition)
        else:
            return jnp.where(condition, x, y)

    def copy(self, a):
        # No need to copy, JAX arrays are immutable