如何将 ndarray 的转换部分分配给自身？

Question

我想要以下 pytorch 代码的行为：

def random_w(L,A):
    w = torch.rand((L,L,A,A))
    # We want (i,i,a,b) = 0
    diag = w.diagonal()[...] = 0
    # And (i,j,a,b) = (j,i,b,a)
    i,j = torch.triu_indices(L,L)
    w[i,j] = w[j,i].transpose(1,2)
    return w

我有以下 Rust 代码（暂时用

ones

替换）

use ndarray::prelude::*;

fn random<F : NdFloat>(L : usize, A: usize) -> Array4<F> {
    let mut w = Array::<F, Ix4>::ones((L,L,A,A));
    for i in 0..L {
        w.slice_mut(s![i,i,..,..]).fill(F::from(0.0).unwrap());
    }
    for i in 0..L {
        for j in (i+1)..L {
            let transpose = w.slice(s![j,i,..,..]).t().into_owned();
            w.slice_mut(s![i,j,..,..]).assign(&transpose);
        }
    }
    w
}

有没有办法避免这里的复制？我发现我通常无法安全地从视图分配到数组中，但在我的情况下，两个区域都不重叠。

Answer 1

我解决了

use ndarray::prelude::*;

fn random<F : NdFloat>(L : usize, A: usize) -> Array4<F> {
    let mut w = Array::<F, Ix4>::ones((L,L,A,A));
    for i in 0..L {
        w.slice_mut(s![i,i,..,..]).fill(F::from(0.0).unwrap());
    }
    for i in 0..L {
        for j in (i+1)..L {
            for a in 0..A {
                for b in 0..A {
                    w[(j,i,b,a)] = w[(i,j,a,b)];
                }
            }
        }
    }
    w
}

不太理想，但由于 Rust 已编译，我想性能会很好