numpy_helper.py 文件源码

def get_chunks(x, nbytes_desired):
    nbytes = np.array(x).ravel()[0].nbytes
    size_desired = nbytes_desired / nbytes
    if size_desired >= x.size:
        # desired chunk size is greater or equal than array size, thus we can include the whole array in a single chunk
        return x.shape
    s = x.shape[::-1]
    cp = np.cumprod(s)
    dim = np.argmax(cp >= size_desired)
    s_dim_desired = size_desired / np.prod(s[:dim])
    s_dim = np.round(s_dim_desired)
    if s_dim < 1:
        s_dim = 1
    chunks = np.ones_like(s)
    chunks[:dim] = s[:dim]
    chunks[dim] = s_dim
    result = tuple(chunks[::-1])
    return result