def promote_rank_to_same_as_data(
    param: np.ndarray, data: np.ndarray, axis: Optional[int]
) -> np.ndarray:
    """
    Promote param (i.e. zero point or scale) rank to same as quantized data,
    so subtraction or multiplication can happen properly on the specified axis.

    For example, when `param` has shape [2] while `data` has shape [2, 2, 2],
    - if `axis` is 0, `param` will be promoted to shape [2, 1, 1].
    - if `axis` is 1, `param` will be promoted to shape [1, 2, 1].
    - if `axis` is 2, `param` will be promoted to shape [1, 1, 2].
    """
    if len(param.shape) == len(data.shape):
        return param  # already have the same rank

    if axis is not None:
        axis = axis if axis >= 0 else axis + len(data.shape)

    if len(param.shape) == 0:
        return np.reshape(param, np.ones(len(data.shape), np.int32))
    else:
        axes = [i for i in range(len(data.shape)) if i != axis]
        return np.expand_dims(param, axis=tuple(axes))

    Promote param (i.e. zero point or scale) rank to same as quantized data,
    so subtraction or multiplication can happen properly on the specified axis.

    For example, when `param` has shape [2] while `data` has shape [2, 2, 2],
    - if `axis` is 0, `param` will be promoted to shape [2, 1, 1].
    - if `axis` is 1, `param` will be promoted to shape [1, 2, 1].
    - if `axis` is 2, `param` will be promoted to shape [1, 1, 2].
    
coremltools/optimize/_utils.py