def where(condition, if_true, if_false):
  """Selects slices of `if_true` and `if_false` based on `condition`.

  This is a wrapper around tf.where() that allows it to select slices at any
  level of a tensor.  (Recall that tf.where() can only select slices on the
  first axis or select individual elements).

  Example usage:

  condition = tf.constant([[True, False],
                           [False, True]])
  if_true = tf.constant([[[1, 1], [2, 2]],
                         [[3, 3], [4, 4]]])
  if_false = tf.constant([[[5, 5], [6, 6]],
                          [[7, 7], [8, 8]]])

  result = where(condition, if_true, if_false)

  assert result.eval() == [[[1, 1], [6, 6]],
                           [[7, 7], [4, 4]]]

  Args:
    condition: <bool>[...] Condition for selecting slices.
    if_true: <T>[condition.shape, ...] Values to use if `condition` is true.
    if_false: <T>[condition.shape, ...] Values to use if `condition` is false.

  Returns:
    <T>[if_true.shape] Values after conditional selection.
  """
  # TODO(terrykoo): Rewrite when tf.where() supports broadcasting: b/78594182.
  condition_rank = condition.shape.ndims
  result_rank = if_true.shape.ndims

  # Use tf.where() directly if it can handle the args.
  if condition_rank == 1 or condition_rank == result_rank:
    return tf.where(condition, if_true, if_false)

  result_dims = shape(if_true)
  inner_dims = result_dims[condition_rank:]

  # Flatten outer dims so we can select slices on axis=0.
  condition = tf.reshape(condition, [-1])
  if_true = tf.reshape(if_true, [-1] + inner_dims)
  if_false = tf.reshape(if_false, [-1] + inner_dims)

  result = tf.where(condition, if_true, if_false)

  # Restore original shape.
  return tf.reshape(result, result_dims)
Selects slices of `if_true` and `if_false` based on `condition`.

  This is a wrapper around tf.where() that allows it to select slices at any
  level of a tensor.  (Recall that tf.where() can only select slices on the
  first axis or select individual elements).

  Example usage:

  condition = tf.constant([[True, False],
                           [False, True]])
  if_true = tf.constant([[[1, 1], [2, 2]],
                         [[3, 3], [4, 4]]])
  if_false = tf.constant([[[5, 5], [6, 6]],
                          [[7, 7], [8, 8]]])

  result = where(condition, if_true, if_false)

  assert result.eval() == [[[1, 1], [6, 6]],
                           [[7, 7], [4, 4]]]

  Args:
    condition: <bool>[...] Condition for selecting slices.
    if_true: <T>[condition.shape, ...] Values to use if `condition` is true.
    if_false: <T>[condition.shape, ...] Values to use if `condition` is false.

  Returns:
    <T>[if_true.shape] Values after conditional selection.
  
language/common/utils/tensor_utils.py