Question

我想广播到子阵列（即广播到数组切片）。这在GPU编程中很有用，例如我想要：

X,Y,Z = (rand(3,3,3) for _=1:3)
@.[1,2] X = f(2X^2 + 6X^3 - sqrt(X)) + Y*Z

其中@.[1,2]表示沿着暗淡3广播，即将冒号应用于表达式中的暗淡1和2。

有没有办法支持这种“子广播”？

编辑：添加示例

julia> a = reshape(1:8, (2,2,2))
2×2×2 Base.ReshapedArray{Int64,3,UnitRange{Int64},Tuple{}}:
[:, :, 1] =
 1  3
 2  4

[:, :, 2] =
 5  7
 6  8

julia> broadcast(*, a, a)
2×2×2 Array{Int64,3}:
[:, :, 1] =
 1   9
 4  16

[:, :, 2] =
 25  49
 36  64

julia> broadcast(*, a, a, dim=3) # I would like to broadcast the matrix multiplication (batch mode) instead of elementwise multiplication.
2×2×2 Array{Int64,3}:
[:, :, 1] =
  7  15
 10  22

[:, :, 2] =
 67   91
 78  106

编辑2：我在这里通过ArrayFire.jl包（arrayfire的包装器）尝试不同的矢量化方法https://arrayfire.com/introduction-to-vectorization/，包括矢量化，并行for循环，批处理和高级矢量化。 arrayfire具有gfor（http://arrayfire.org/docs/page_gfor.htm）方法，用于在矩阵切片上运行并行计算，并通过ArrayFire.jl中的广播实现。目前，朱莉娅的广播是元素化的。我只是想知道它是否能够“切片”，然后它可以为线性代数函数（https://github.com/arrayfire/arrayfire/issues/483）提供纯julia 3D和4D支持。

当然，正常的嵌套for循环将完成工作。我刚刚退出广播.语法，并想知道它是否可以扩展。

Answer 1

我认为您正在寻找mapslices。

mapslices(x->x*x, a, (1,2))
2×2×2 Array{Int64,3}:
[:, :, 1] =
  7  15
 10  22

[:, :, 2] =
 67   91
 78  106

mapslices(f, A, dims)

使用函数f转换数组A的给定维数。 f是   调用A [...，：，...，：，...]形式的A的每个切片。 dims是一个   整数向量，指定冒号在此表达式中的位置。
  结果沿其余维度连接。对于   例如，如果dims为[1,2]且A为4维，则调用f   所有i和j都有[：，：，i，j]。

如果要指定要连接的维度而不是应用该函数的维度，请使用setdiff。

（如果您需要多参数版本，请查看此要点https://gist.github.com/alexmorley/e585df0d8d857d7c9e4a5af75df43d00）

可以广播应用于julia中的子阵列/数组切片

1 个答案: