easy: implement mtf.expand_dims and mtf.squeeze about mesh HOT 1 OPEN

def squeeze(x, dim_name, name=None):
    # TODO: assert the dimension appears in x with size 1?
    return mtf.reduce_sum(x, reduced_dim=mtf.Dimension(dim_name, 1), name=name)

def expand_dims(x, dim_name, axis, name=None):
    new_dims = list(x.shape[:])
    if axis == -1:
        new_dims.append(mtf.Dimension(dim_name, 1))
    elif axis < 0:
        new_dims.insert(axis + 1, mtf.Dimension(dim_name, 1))
    else:
        new_dims.insert(axis, mtf.Dimension(dim_name, 1))
    return mtf.reshape(x, mtf.Shape(new_dims), name=name)