CODEWITHSUNDEEP
rustmultidimensional-arrayrust-ndarray
Unlikus
Unlikus

Reputation: 1858

How do I assign a transformed part of an ndarray to itself?

I want the behavior of the following pytorch code:

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

I have the following Rust code (random replaced with ones for now)

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
}

Is there a way to avoid the copy here? I see that I cannot safely assign from a View into the array in general, but in my case both regions are non overlapping.

Upvotes: 0

Views: 63

Answers (1)

Unlikus
Unlikus

Reputation: 1858

I settled with

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
}

Not ideal, but since Rust is compiled, I guess the performance will be fine

Upvotes: 0

Related Questions