def vectorize(
    current_tensor, cluster_dim, vector_ch_axis=0
) -> _Tuple[_torch.Tensor, _torch.Tensor]:
    """
    Function to vectorize a tensor till the point where its numel is divisible by cluster_dim. The remaining parameters
    are returned as a pad.

    vector_ch_axis 0
    tensor with shape [x_1, x_2, x_3, ..., x_N] --> [x_N, x_2, x_3, ..., x_(N-1), x_1] --> [x_N*x_2*x_3* ... *x_(N-1)*x_1/cluster_dim, cluster_dim]

    vector_ch_axis 1
    tensor with shape [x_1, x_2, x_3, ..., x_N] --> [x_1, x_N, x_3, ..., x_(N-1), x_2] --> [x_1*x_N*x_3* ... *x_(N-1)*x_2/cluster_dim, cluster_dim]
    """
    num_misalignment = _torch.numel(current_tensor) % cluster_dim

    if cluster_dim > 1:
        current_tensor = current_tensor.transpose(vector_ch_axis, -1)

    pad = None
    if num_misalignment:
        current_tensor = current_tensor.flatten()
        pad = current_tensor[-num_misalignment:]
        current_tensor = current_tensor[:-num_misalignment]

    return current_tensor.reshape(-1, cluster_dim), pad

    Function to vectorize a tensor till the point where its numel is divisible by cluster_dim. The remaining parameters
    are returned as a pad.

    vector_ch_axis 0
    tensor with shape [x_1, x_2, x_3, ..., x_N] --> [x_N, x_2, x_3, ..., x_(N-1), x_1] --> [x_N*x_2*x_3* ... *x_(N-1)*x_1/cluster_dim, cluster_dim]

    vector_ch_axis 1
    tensor with shape [x_1, x_2, x_3, ..., x_N] --> [x_1, x_N, x_3, ..., x_(N-1), x_2] --> [x_1*x_N*x_3* ... *x_(N-1)*x_2/cluster_dim, cluster_dim]
    
coremltools/optimize/torch/palettization/_utils.py