def METHOD_NAME(shape: List[int]) -> List[int]:
    """
    Utility method to ensure that 1-dimensional shapes
    assume that `[N] == [N, D]`.
    """
    if len(shape) > 1:
        return shape
    return shape + [1]