def METHOD_NAME(op, **kwargs):
    offset = op.offset
    # axis1 = op.axis1
    # axis2 = op.axis2
    @numba_basic.numba_njit(inline="always")
    def extract_diag(x):
        return np.diag(x, k=offset)
    return extract_diag