Make x n-d with squeeze and expand_dims.

def expand_squeeze_to_nd(x, n, squeeze_dim=2, expand_dim=-1):
  """Make x n-d with squeeze and expand_dims."""
  if len(x.shape) > n:
    while len(x.shape) != n:
      x = tf.squeeze(x, [squeeze_dim])
  else:
    while len(x.shape) != n:
      x = tf.expand_dims(x, expand_dim)
  return x