def reshape_weight_for_vector_lut(
    weight: np.ndarray, vector_size: int, vector_axis: int
) -> np.ndarray:
    """
    For vector palettization, we need to extract vectors and move them to the last dim.
    If the input weight has shape [s0, s1, s2, ... , sn], the output shape should have shape
        [s0, ..., si // vector_size, ..., sn, vector_size] where i == vector_axis.

    For example, starting from weight `a` which has shape `[4, 4]` and `vector_size=2` and
    `vector_axis=0`, we want to reshape the matrix into `c` with shape `[2, 4, 2]`, where `c[0, 0]`
    contains `a[0,0], a[1, 0]`, and `c[1, 0]` contains `a[2, 0], a[3, 0]`, etc.

    To achieve this, we need to first swap the vector_axis to last dim, and split out the vector_size,
    and finally swap it back. Here is a concrete exmaple:
        a = np.array([[ 0,  1,  2,  3],
                      [ 4,  5,  6,  7],
                      [ 8,  9, 10, 11],
                      [12, 13, 14, 15]])
        b = np.swapaxes(a, 0, -1).reshape((4, 2, 2))
        c = np.swapaxes(b, 0, 1)
        # c[0, 0] is array([0, 4])
        # c[1, 0] is array([8, 12])
    """
    weight = np.swapaxes(weight, -1, vector_axis)
    weight = weight.reshape((*weight.shape[:-1], weight.shape[-1] // vector_size, vector_size))
    return np.swapaxes(weight, -2, vector_axis)

    For vector palettization, we need to extract vectors and move them to the last dim.
    If the input weight has shape [s0, s1, s2, ... , sn], the output shape should have shape
        [s0, ..., si // vector_size, ..., sn, vector_size] where i == vector_axis.

    For example, starting from weight `a` which has shape `[4, 4]` and `vector_size=2` and
    `vector_axis=0`, we want to reshape the matrix into `c` with shape `[2, 4, 2]`, where `c[0, 0]`
    contains `a[0,0], a[1, 0]`, and `c[1, 0]` contains `a[2, 0], a[3, 0]`, etc.

    To achieve this, we need to first swap the vector_axis to last dim, and split out the vector_size,
    and finally swap it back. Here is a concrete exmaple:
        a = np.array([[ 0,  1,  2,  3],
                      [ 4,  5,  6,  7],
                      [ 8,  9, 10, 11],
                      [12, 13, 14, 15]])
        b = np.swapaxes(a, 0, -1).reshape((4, 2, 2))
        c = np.swapaxes(b, 0, 1)
        # c[0, 0] is array([0, 4])
        # c[1, 0] is array([8, 12])
    
coremltools/optimize/_utils.py