def _align_dims(tensor: tf.Tensor, target_ndims: int) -> tf.Tensor:
  """Expand the rank of input tensor until it matches the target rank.

  Non-elementwise per-key reduce returns a tensor with rank 1 (batch).
  The dimension count needs to match with x to finish the final mapping, because
  we want to broadcast each reduction with x. To do so we need to add singleton
  dimensions, otherwise TF will try to broadcast along the wrong dimensions.

  Args:
    tensor: A `Tensor`.
    target_ndims: The count of dims we want the output to meet or exceed.

  Returns:
    The original input, with dimension count >= target_ndims.
  """
  if target_ndims is None or target_ndims <= tensor.get_shape().ndims:
    return tensor
  for _ in range(target_ndims - tensor.get_shape().ndims):
    tensor = tf.expand_dims(tensor, -1)
  return tensor
Expand the rank of input tensor until it matches the target rank.

  Non-elementwise per-key reduce returns a tensor with rank 1 (batch).
  The dimension count needs to match with x to finish the final mapping, because
  we want to broadcast each reduction with x. To do so we need to add singleton
  dimensions, otherwise TF will try to broadcast along the wrong dimensions.

  Args:
    tensor: A `Tensor`.
    target_ndims: The count of dims we want the output to meet or exceed.

  Returns:
    The original input, with dimension count >= target_ndims.
  
tensorflow_transform/tf_utils.py