code: def expand_squeeze_to_nd(x, n, squeeze_dim=2, expand_dim=-1):
    
    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
docs: Make x n-d with squeeze and expand_dims.