def where(condition, x=None, y=None):
  """Raises ValueError if exactly one of x or y is not None."""
  condition = asarray(condition, dtype=np.bool_)
  if x is None and y is None:
    return nonzero(condition)
  elif x is not None and y is not None:
    x, y = _promote_dtype(x, y)
    return utils.tensor_to_ndarray(tf.where(condition.data, x.data, y.data))
  raise ValueError('Both x and y must be ndarrays, or both must be None.')
Raises ValueError if exactly one of x or y is not None.
trax/tf_numpy/numpy_impl/array_ops.py