Scatter/Gather 喜欢 ArrayFire 中的 Numpy

Scatter/Gather like Numpy in ArrayFire

我想在一个轴上的特定索引处分散和收集数组 X 中的元素。 因此,给定一个索引数组 idx,我想 select 第 0 列的第 idx(0) 个元素,第 1 列的第 idx(1) 个元素,等等..

在 Numpy 中,如下语句:

X = np.array([[1, 2, 3], [4, 5, 6]])
print(X[[0, 1, 1], range(3)])

打印 [1, 5, 6]

此外,我可以反向执行此过程:

Y = np.zeros((2, 3))
Y[[0, 1, 1], range(3)] = [1, 5, 6]
print(Y)

这将打印

[[1. 0. 0.]
 [0. 5. 6.]]

但是,当我尝试在 ArrayFire 中复制此行为时:

float elements[] = {1, 2, 3, 4, 5, 6};
af::array X = af::array(3, 2, elements);
int idx_elements[] = {0, 1, 1};
af::array idx = af::array(3, idx_elements);
af::print("", X(af::span, idx));

我得到一个形状为 [3, 3, 1, 1] 的数组,元素为

1.0000     4.0000     4.0000 
2.0000     5.0000     5.0000 
3.0000     6.0000     6.0000

那么我怎样才能在 ArrayFire 中实现分散和聚集元素所需的类似 numpy 的行为?

要对矩阵执行聚集操作,我可以提取结果矩阵的对角线,但这在多维情况下可能不起作用,在另一个(分散)方向上也不起作用。

X
[3 2 1 1]
    1.0000     4.0000 
    2.0000     5.0000 
    3.0000     6.0000 

idx
[3 1 1 1]
         0 
         1 
         1 

ArrayFire 在涉及 af::array 时计算笛卡尔积。因此,输出。 因此,请查看以下指数。

Col\Row   0       1        1   from array

   0    (0, 0)  (0,1)  (0, 1)
   1    (1, 0)  (1,1)  (1, 1)
   2    (2, 0)  (2,1)  (2, 1)
   ^
   ^ from sequence

因此,X(af::span, idx)) 的输出是一个 3x3 矩阵。

要根据坐标收集元素,您需要不同的函数 approx2。请注意,此函数仅将其索引作为浮点数组。

float idx_elements[] = {0, 1, 1}; // changed the idx to floats
af::array colIdx = af::array(3, idx_elements);
af::array rowIdx = af::iota(3); // same effect as span

af::array out = approx2(X, rowIdx, colIdx);
af_print(out); 
// out
// [3 1 1 1]
//     1.0000 
//     5.0000 
//     6.0000 

要设置给定索引的值,出于某种原因,您必须展平数组 当涉及 af::array 时,array::operator() 考虑笛卡尔积。

af::array A = af::constant(0, 3, 2); // same size as X
af::array B = af::flat(A); // flatten the array, this involves meta data modification only

B(rowIdx + 3 * colIdx) = out; // use row & col indices to fetch linear indices
// rowIdx + 3 * colIdx
// [3 1 1 1]
//     0.0000 
//     4.0000 
//     5.0000 

B = moddims(B, A.dims()); // reset the dimensions to original A dims

af_print(B);
// B
// [3 2 1 1]
//     1.0000     0.0000 
//     0.0000     5.0000 
//     0.0000     6.0000 

您可以在我们的 indexing tutorial.

中查看更多详细信息