build diagonal batched tensor with unknown dimensions

[japm48]

I'm trying to implement a simplified version of tf.linalg.diag. Basically, the input is an n-dimensional tensor with shape [..., k], then the output is (n+1)-dimensional with shape [..., k, k].
The output elements are:

• output[..., i, j] = 0: if i!=j
• output[..., i, j] = input[..., i]: if i==j

Of course, everything is written in terms of dynamic/IxDyn dimensions.
My idea was to start zero-initializing the output matrix and then fill in the non-zero elements, but then I cannot assign to a sliced view as the indices of the diagonal are varying (that would be suitable instead for filling a row or a column).

I'm a loss about what to do. Is there an "idiomatic" approach to solve this kind of problem?

[daniellga]

I don't know if I understood correctly what it does, but can you check if this works? It probably lacks a lot of checks but the basic stuff is there I think. I also didn't manage to use exact_chunks with Zip and had to use zip instead... It looks like it doesn't work if matrix and diag don't have the same shape.

fn transform_dim(tuple: &[usize], m: usize) -> Vec<usize> {
    let n = tuple.len();
    tuple
        .iter()
        .enumerate()
        .map(|(i, &val)| if i < n - m { 1 } else { val })
        .collect()
}

#[cfg(test)]
mod tests {
    use ndarray::{ArrayBase, IxDyn};

    use crate::transform_dim;

    #[test]
    fn test_diag1() {
        use ndarray::Array;

        let diagonal: ArrayBase<_, IxDyn> =
            Array::from_shape_vec(IxDyn(&[2, 4]), vec![1, 2, 3, 4, 5, 6, 7, 8])
                .expect("Failed to create the array")
                .into_dyn();

        let mut new_shape = diagonal.shape().to_vec();
        new_shape.push(*new_shape.last().unwrap());
        let new_dimension = IxDyn(&new_shape);

        let mut result = Array::<i32, IxDyn>::zeros(new_dimension);

        let shape1 = transform_dim(result.shape(), 2);
        let shape2 = transform_dim(diagonal.shape(), 1);

        for (matrix, diag) in result
            .exact_chunks_mut(IxDyn(&shape1))
            .into_iter()
            .zip(diagonal.exact_chunks(IxDyn(&shape2)))
        {
            {
                let mut matrix = matrix.squeeze();
                let diag = diag.squeeze();
                let mut d = matrix.diag_mut();
                d.assign(&diag);
            }
        }

        println!("{:?}", result);
    }

    #[test]
    fn test_diag2() {
        use ndarray::Array;

        let diagonal: ArrayBase<_, IxDyn> = Array::from_shape_vec(
            IxDyn(&[2, 4, 4]),
            vec![
                1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23,
                24, 25, 26, 27, 28, 29, 30, 31, 32,
            ],
        )
        .expect("Failed to create the array")
        .into_dyn();

        let mut new_shape = diagonal.shape().to_vec();
        new_shape.push(*new_shape.last().unwrap());
        let new_dimension = IxDyn(&new_shape);

        let mut result = Array::<i32, IxDyn>::zeros(new_dimension);

        let shape1 = transform_dim(result.shape(), 2);
        let shape2 = transform_dim(diagonal.shape(), 1);

        for (matrix, diag) in result
            .exact_chunks_mut(IxDyn(&shape1))
            .into_iter()
            .zip(diagonal.exact_chunks(IxDyn(&shape2)))
        {
            {
                let mut matrix = matrix.squeeze();
                let diag = diag.squeeze();
                let mut d = matrix.diag_mut();
                let b = diag.flatten();
                d.assign(&b);
            }
        }

        println!("{:?}", result);
    }
}

[japm48]

Sorry for the delay, I 100% forgot about this.
Thanks a lot! I ended up opting for another alternative, but I believe your solution does the trick.