dask.array.blockwise

dask.array.blockwise

dask.array.blockwise(func, out_ind, *args, name=None, token=None, dtype=None, adjust_chunks=None, new_axes=None, align_arrays=True, concatenate=None, meta=None, **kwargs)[源代码]

张量操作:广义内积和外积

一类广泛的阻塞算法和模式可以用简洁的多索引表示法来指定。blockwise 函数以多种方式将内存中的函数应用于多个输入的多个块。许多 dask.array 操作是 blockwise 的特例,包括元素级操作、广播、归约、张量点积和转置。

参数
函数可调用

应用于块的单个元组的函数

out_ind可迭代对象

输出块的模式,类似于 ‘ijk’ 或 (1, 2, 3)

*args数组序列,索引对

您也可以传递字面参数,伴随 None 索引,例如 (x, ‘ij’, y, ‘jk’, z, ‘i’, some_literal, None)

**kwargsdict

传递给函数的额外关键字参数

dtypenp.dtype

结果数组的类型。

连接bool, 仅关键字

如果为真,则沿虚拟索引连接数组,否则提供列表

调整块dict

字典映射索引到应用于块大小的函数

新轴dict, 仅关键字

新的索引及其维度长度

align_arrays: bool

当提供多个数组时,是否沿等尺寸的维度对齐块。这允许某些数组中的较大块被分解为与其它数组中的块大小匹配的较小块,以便它们能够兼容块函数映射。如果这是false,那么如果数组在每个维度中没有相同数量的块,则会抛出错误。

示例

从两个数组 x 和 y 进行的 2D 非常规并行操作。

>>> import operator, numpy as np, dask.array as da
>>> x = da.from_array([[1, 2],
...                    [3, 4]], chunks=(1, 2))
>>> y = da.from_array([[10, 20],
...                    [0, 0]])
>>> z = blockwise(operator.add, 'ij', x, 'ij', y, 'ij', dtype='f8')
>>> z.compute()
array([[11, 22],
       [ 3,  4]])

外积将两个一维向量 a 和 b 相乘

>>> a = da.from_array([0, 1, 2], chunks=1)
>>> b = da.from_array([10, 50, 100], chunks=1)
>>> z = blockwise(np.outer, 'ij', a, 'i', b, 'j', dtype='f8')
>>> z.compute()
array([[  0,   0,   0],
       [ 10,  50, 100],
       [ 20, 100, 200]])

z = x.T

>>> z = blockwise(np.transpose, 'ji', x, 'ij', dtype=x.dtype)
>>> z.compute()
array([[1, 3],
       [2, 4]])

上面的转置案例具有说明性,因为它通过调用 np.transpose 对每个内存块进行了转置,并通过交换索引顺序 ij -> ji 对块本身的顺序进行了转置。

我们可以用更多的变量和更复杂的内存函数来组合这些相同的模式

z = X + Y.T

>>> z = blockwise(lambda x, y: x + y.T, 'ij', x, 'ij', y, 'ji', dtype='f8')
>>> z.compute()
array([[11,  2],
       [23,  4]])

任何在输出索引中缺失的索引,如 i ,都被解释为缩并(注意这与爱因斯坦约定不同;重复的索引并不意味着缩并。) 在缩并的情况下,传递的函数应期望在任何包含该索引的数组上接收一个块的可迭代对象。 若要沿着缩并的维度接收连接的数组,请传递 concatenate=True

内积乘以 a 和 b,两个一维向量

>>> def sequence_dot(a_blocks, b_blocks):
...     result = 0
...     for a, b in zip(a_blocks, b_blocks):
...         result += a.dot(b)
...     return result
>>> z = blockwise(sequence_dot, '', a, 'i', b, 'i', dtype='f8')
>>> z.compute()
np.int64(250)

使用 new_axes= 关键字添加新的单块维度,包括新维度的长度。新维度将始终位于单个块中。

>>> def f(a):
...     return a[:, None] * np.ones((1, 5))
>>> z = blockwise(f, 'az', a, 'a', new_axes={'z': 5}, dtype=a.dtype)

新的维度也可以通过指定一个块大小的元组来实现多块。这本身具有有限的实用性(因为块都是相同的),但生成的图可以被修改以实现更有用的结果(参见 da.map_blocks)。

>>> z = blockwise(f, 'az', a, 'a', new_axes={'z': (5, 5)}, dtype=x.dtype)
>>> z.chunks
((1, 1, 1), (5, 5))

如果应用的函数改变了每个块的大小,你可以通过一个 adjust_chunks={...} 字典来指定,该字典为每个索引持有一个函数,用于修改该索引中的维度大小。

>>> def double(x):
...     return np.concatenate([x, x])
>>> y = blockwise(double, 'ij', x, 'ij',
...               adjust_chunks={'i': lambda n: 2 * n}, dtype=x.dtype)
>>> y.chunks
((2, 2), (2,))

通过使用 None 进行索引,包含文字字面量

>>> z = blockwise(operator.add, 'ij', x, 'ij', 1234, None, dtype=x.dtype)
>>> z.compute()
array([[1235, 1236],
       [1237, 1238]])