def reduce_batch_minus_min_and_max_per_key(
    x: common_types.TensorType,
    key: common_types.TensorType,
    reduce_instance_dims: bool = True
) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
  """Computes the -min and max of a tensor x.

  Args:
    x: A `Tensor` or `CompositeTensor`.
    key: A `Tensor` or `CompositeTensor`. Must meet one of the following
      conditions - 1. Both x and key are dense, 2. Both x and key are composite
      and `key` must exactly match `x` in everything except values, 3. The
      axis=1 index of each element of sparse x matches its index of dense key.
    reduce_instance_dims: A bool indicating whether this should collapse the
      batch and instance dimensions to arrive at a single scalar output, or only
      collapse the batch dimension and outputs a vector of the same shape as the
      input.
  Returns:
    A 3-tuple containing the `Tensor`s (key_vocab, min_per_key, max_per_key).
  """
  if x.dtype == tf.uint8 or x.dtype == tf.uint16:
    x = tf.cast(x, tf.int32)

  elif x.dtype == tf.uint32 or x.dtype == tf.uint64:
    raise TypeError('Tensor type %r is not supported' % x.dtype)

  if not reduce_instance_dims and isinstance(
      x, (tf.SparseTensor, tf.RaggedTensor)):
    raise NotImplementedError(
        'Elementwise reduction of composite tensors is not supported'
    )

  x, key = _validate_and_get_dense_value_key_inputs(x, key)

  def get_batch_max_per_key(tensor, key_uniques):  # pylint: disable=missing-docstring
    if not reduce_instance_dims or tensor.get_shape().ndims < 2:
      row_maxes = tensor
    else:
      row_maxes = tf.reduce_max(
          tensor, axis=tf.range(1, tensor.get_shape().ndims))
    return tf.math.unsorted_segment_max(row_maxes, key_uniques.idx,
                                        tf.size(input=key_uniques.y))

  unique = tf.unique_with_counts(key, out_idx=tf.int64)
  x_batch_maxes = get_batch_max_per_key(x, unique)
  x_batch_minus_mins = get_batch_max_per_key(-x, unique)

  return (unique.y, x_batch_minus_mins, x_batch_maxes)
Computes the -min and max of a tensor x.

  Args:
    x: A `Tensor` or `CompositeTensor`.
    key: A `Tensor` or `CompositeTensor`. Must meet one of the following
      conditions - 1. Both x and key are dense, 2. Both x and key are composite
      and `key` must exactly match `x` in everything except values, 3. The
      axis=1 index of each element of sparse x matches its index of dense key.
    reduce_instance_dims: A bool indicating whether this should collapse the
      batch and instance dimensions to arrive at a single scalar output, or only
      collapse the batch dimension and outputs a vector of the same shape as the
      input.
  Returns:
    A 3-tuple containing the `Tensor`s (key_vocab, min_per_key, max_per_key).
  
tensorflow_transform/tf_utils.py