def _broadcast_to_x_shape(x, y):
  """Broadcasts y to same shape as x as needed.

  Args:
    x: An input feature.
    y: A feature that is either the same shape as x or has the same outer
      dimensions as x. If the latter, y is broadcast to the same shape as x.

  Returns:
    A Tensor that contains the broadcasted feature, y.
  """
  # The batch dimension of x and y must be the same, and y must be 1D.
  x_shape = tf.shape(input=x)
  y_shape = tf.shape(input=y)
  assert_eq = tf.compat.v1.assert_equal(x_shape[0], y_shape[0])
  with tf.control_dependencies([assert_eq]):
    rank_delta = tf.rank(x) - tf.rank(y)
  target_shape = tf.concat(
      [tf.shape(y), tf.ones(rank_delta, dtype=tf.int32)], axis=0)
  matched_rank = tf.reshape(y, target_shape)
  return tf.broadcast_to(matched_rank, x_shape)
Broadcasts y to same shape as x as needed.

  Args:
    x: An input feature.
    y: A feature that is either the same shape as x or has the same outer
      dimensions as x. If the latter, y is broadcast to the same shape as x.

  Returns:
    A Tensor that contains the broadcasted feature, y.
  
tensorflow_transform/tf_utils.py