Repository: webonnx/wonnx (language: python)

Issue:
Incorrect results for Transpose
**Describe the bug**
Transpose appears to produce incorrect results for certain permutations.

**To Reproduce**

In the following test case, the following work (like in NumPy):
* Transpose perm=`0,2,1,3` followed by Transpose perm=`0,2,1,3`which should reverse the first transpose
* Transpose perm=`0,3,2,1` followed by Transpose perm=`0,3,2,1`which should reverse the first transpose

The following works in NumPy (see expected behaviour below), but fails in woonx:
* Transpose perm=`0,2,3,1` followed by Transpose perm=`0,3,1,2` which should reverse the first transpose

(Note that in this case, the first perm is not equal to the latter).

````rust

fn test_transpose_4d_perm(transpose_first: &[i64], transpose_second: &[i64]) {
    let mut input_data = HashMap::new();
    let data = (0..2 * 3 * 4).map(|x| x as f32).collect::<Vec<f32>>();
    input_data.insert("X".to_string(), data.as_slice().into());

    let x_dims = vec![1, 2, 3, 4];
    let intermediate_dims: Vec<i64> = transpose_first
        .iter()
        .map(|i| x_dims[*i as usize])
        .collect();

    // Model: X -> Transpose -> Y -> Transpose -> Z; X==Z
    let model = model(graph(
        vec![tensor("X", &x_dims)],
        vec![tensor("Z", &x_dims)],
        vec![tensor("Y", &intermediate_dims)],
        vec![],
        vec![
            node(
                vec!["X"],
                vec!["Y"],
                "Transpose",
                "Transpose",
                vec![attribute("perm", transpose_first.to_vec())],
            ),
            node(
                vec!["Y"],
                vec!["Z"],
                "Transpose",
                "Transpose",
                vec![attribute("perm", transpose_second.to_vec())],
            ),
        ],
    ));

    let session =
        pollster::block_on(wonnx::Session::from_model(model)).expect("session did not create");
    let result = pollster::block_on(session.run(&input_data)).unwrap();

    common::assert_eq_vector((&result["Z"]).try_into().unwrap(), &data);
}

/* This tests the equivalent of the following Python code:
a = np.arange(0,24).reshape((1,2,3,4));
a == a.transpose(a).transpose(inverse of a)
*/
#[test]
fn test_two_transposes_4d() {
    // a == a.transpose([0,2,1,3]).transpose([0,2,1,3])
    test_transpose_4d_perm(&[0, 2, 1, 3], &[0, 2, 1, 3]);

    // ! WORKS in python, FAILS in wonnx...
    // a == a.transpose([0,2,3,1]).transpose([0,3,1,2])
    // test_transpose_4d_perm(&[0, 2, 3, 1], &[0, 3, 1, 2]);
    // a == a.transpose([0,3,2,1]).transpose([0,3,2,1])
    test_transpose_4d_perm(&[0, 3, 2, 1], &[0, 3, 2, 1]);
}
````

**Expected behavior**


````python
>>> a = np.arange(0,24).reshape((1,2,3,4));
>>> a
array([[[[ 0,  1,  2,  3],
         [ 4,  5,  6,  7],
         [ 8,  9, 10, 11]],

        [[12, 13, 14, 15],
         [16, 17, 18, 19],
         [20, 21, 22, 23]]]])
>>> a.transpose([0,2,3,1]).transpose([0,3,1,2])
array([[[[ 0,  1,  2,  3],
         [ 4,  5,  6,  7],
         [ 8,  9, 10, 11]],

        [[12, 13, 14, 15],
         [16, 17, 18, 19],
         [20, 21, 22, 23]]]])
>>> a == a.transpose([0,2,3,1]).transpose([0,3,1,2])
array([[[[ True,  True,  True,  True],
         [ True,  True,  True,  True],
         [ True,  True,  True,  True]],

        [[ True,  True,  True,  True],
         [ True,  True,  True,  True],
         [ True,  True,  True,  True]]]])
````


Discussion / hints:
Not that this would affect your example, but the default permutation is also computed incorrectly:

https://github.com/webonnx/wonnx/blob/61c8b69d2f6e7bc3c672d562b0fbfe93f1225e3e/wonnx/src/compiler.rs#L1288
(it should be something closer to `(0..(input_shapes[0].rank() - 1)).rev().collect()`)

Fix patch:
diff --git a/wonnx/src/compiler.rs b/wonnx/src/compiler.rs
index 27cd99f2..b15d5daa 100644
--- a/wonnx/src/compiler.rs
+++ b/wonnx/src/compiler.rs
@@ -1285,18 +1285,27 @@ pub fn compile(
             }
         }
         "Transpose" => {
-            let default = ((input_lengths[0] as i64)..0).collect::<Vec<_>>();
+            let n_dims: i64 = input_shapes[0].rank() as i64;
+            let default = (0..n_dims).rev().collect::<Vec<i64>>();
             let perms: Vec<i64> = get_attribute("perm", Some(default), node)?;
-            let permuted_shapes = perms
-                .iter()
-                .map(|p| output_shapes[0].dim(*p as usize))
-                .collect::<Vec<_>>();
 
-            let mut chunks = vec![];
-            for i in 1..permuted_shapes.len() {
-                chunks.push(permuted_shapes[i..].iter().product::<u64>());
+            // The number of elements in the permutations list must be equal to the output shape rank
+            if perms.len() != output_shapes[0].rank() {
+                return Err(CompileError::InvalidAttributeValue {
+                    attribute: "perm".to_string(),
+                    value: format!("{:?}", perms),
+                    opset_version,
+                });
             }
-            chunks.push(1);
+
+            let chunks = perms
+                .iter()
+                .map(|p| {
+                    input_shapes[0].dims[((*p as usize) + 1)..]
+                        .iter()
+                        .product::<u64>()
+                })
+                .collect::<Vec<_>>();
 
             context.insert("permuted_chunks", &chunks);
 
diff --git a/wonnx/templates/matrix/transpose.wgsl b/wonnx/templates/matrix/transpose.wgsl
index 98447db3..f55f1675 100644
--- a/wonnx/templates/matrix/transpose.wgsl
+++ b/wonnx/templates/matrix/transpose.wgsl
@@ -10,10 +10,10 @@ var<storage, write> output_0: Array;
 fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
 	let gidx = global_id.x;
 
-	if (gidx < {{ i_lens[0] }}u) {
+	if (gidx < {{ o_lens[0] }}u) {
 		var rest = gidx;
 
-		{%- for chunks in i_chunks[0] -%}
+		{%- for chunks in o_chunks[0] -%}
 			{% if loop.last %}
 				let d_{{ loop.index0 }} = rest; 
 			{% else %}
@@ -32,6 +32,6 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
 			{%- endfor %}
 		;
 
-		output_0.data[index] = input_0.data[gidx];
+		output_0.data[gidx] = input_0.data[index];
 	}
 }
\ No newline at end of file


Test patch:
diff --git a/wonnx-py/tests/test_onnx_backend.py b/wonnx-py/tests/test_onnx_backend.py
index 84b315a4..96d9d154 100644
--- a/wonnx-py/tests/test_onnx_backend.py
+++ b/wonnx-py/tests/test_onnx_backend.py
@@ -154,6 +154,7 @@ def do_enforce_test_coverage_safelist(model):  # type: (ModelProto) -> bool
 backend_test.include(f"test_add_bcast_[a-z,_]*")
 backend_test.include(f"test_sub_bcast_[a-z,_]*")
 backend_test.include(f"test_pow_bcast_[a-z,_]*")
+backend_test.include(f"test_transpose[a-z,_]*")
 
 # Don't support 'bool' type
 # backend_test.include(f"test_and_bcast[a-z0-9,_]*")
diff --git a/wonnx/tests/matrix.rs b/wonnx/tests/matrix.rs
index 80a37bb2..59b0167a 100644
--- a/wonnx/tests/matrix.rs
+++ b/wonnx/tests/matrix.rs
@@ -41,6 +41,132 @@ fn test_matmul_square_matrix() {
     common::assert_eq_vector((&result["C"]).try_into().unwrap(), sum.as_slice().unwrap());
 }
 
+fn test_transpose_4d_perm(transpose_first: &[i64], transpose_second: &[i64]) {
+    let mut input_data = HashMap::new();
+    let data = (0..2 * 3 * 4).map(|x| x as f32).collect::<Vec<f32>>();
+    input_data.insert("X".to_string(), data.as_slice().into());
+
+    let x_dims = vec![1, 2, 3, 4];
+    let intermediate_dims: Vec<i64> = transpose_first
+        .iter()
+        .map(|i| x_dims[*i as usize])
+        .collect();
+
+    // Model: X -> Transpose -> Y -> Transpose -> Z; X==Z
+    let model = model(graph(
+        vec![tensor("X", &x_dims)],
+        vec![tensor("Z", &x_dims)],
+        vec![tensor("Y", &intermediate_dims)],
+        vec![],
+        vec![
+            node(
+                vec!["X"],
+                vec!["Y"],
+                "Transpose",
+                "Transpose",
+                vec![attribute("perm", transpose_first.to_vec())],
+            ),
+            node(
+                vec!["Y"],
+                vec!["Z"],
+                "Transpose",
+                "Transpose",
+                vec![attribute("perm", transpose_second.to_vec())],
+            ),
+        ],
+    ));
+
+    let session =
+        pollster::block_on(wonnx::Session::from_model(model)).expect("session did not create");
+    let result = pollster::block_on(session.run(&input_data)).unwrap();
+
+    common::assert_eq_vector((&result["Z"]).try_into().unwrap(), &data);
+}
+
+/* This tests the equivalent of the following Python code:
+a = np.arange(0,24).reshape((1,2,3,4));
+a == a.transpose(a).transpose(inverse of a)
+*/
+#[test]
+fn test_two_transposes_4d_0213() {
+    // a == a.transpose([0,2,1,3]).transpose([0,2,1,3])
+    test_transpose_4d_perm(&[0, 2, 1, 3], &[0, 2, 1, 3]);
+}
+
+#[test]
+fn test_two_transposes_4d_0312() {
+    // a == a.transpose([0,2,3,1]).transpose([0,3,1,2])
+    test_transpose_4d_perm(&[0, 2, 3, 1], &[0, 3, 1, 2]);
+}
+
+#[test]
+fn test_two_transposes_4d_0321() {
+    // a == a.transpose([0,3,2,1]).transpose([0,3,2,1])
+    test_transpose_4d_perm(&[0, 3, 2, 1], &[0, 3, 2, 1]);
+}
+
+//  a = np.arange(2*3*4).reshape((1,2,3,4))
+// a.transpose((0,3,1,2)).reshape((24,))
+#[test]
+fn test_transposes_4d_0312() {
+    let mut input_data = HashMap::new();
+    let data = (0..2 * 3 * 4).map(|x| x as f32).collect::<Vec<f32>>();
+    input_data.insert("X".to_string(), data.as_slice().into());
+
+    // Model: X -> Transpose -> Y
+    let model = model(graph(
+        vec![tensor("X", &[1, 2, 3, 4])],
+        vec![tensor("Y", &[1, 4, 2, 3])],
+        vec![],
+        vec![],
+        vec![node(
+            vec!["X"],
+            vec!["Y"],
+            "Transpose",
+            "Transpose",
+            vec![attribute("perm", vec![0, 3, 1, 2])],
+        )],
+    ));
+
+    let session =
+        pollster::block_on(wonnx::Session::from_model(model)).expect("session did not create");
+    let result = pollster::block_on(session.run(&input_data)).unwrap();
+
+    common::assert_eq_vector(
+        (&result["Y"]).try_into().unwrap(),
+        &[
+            0., 4., 8., 12., 16., 20., 1., 5., 9., 13., 17., 21., 2., 6., 10., 14., 18., 22., 3.,
+            7., 11., 15., 19., 23.,
+        ],
+    );
+}
+
+// When no 'perm' attribute is specified, Transpose should reverse axes (e.g. for 4D input, assume perm is [3,2,1,0])
+#[test]
+fn test_two_transposes_default_4d() {
+    let mut input_data = HashMap::new();
+    let data = (0..2 * 3 * 4).map(|x| x as f32).collect::<Vec<f32>>();
+    input_data.insert("X".to_string(), data.as_slice().into());
+
+    // Model: X -> Transpose -> Y -> Transpose -> Z; X==Z
+    let model = model(graph(
+        vec![tensor("X", &[2, 3, 4])],
+        vec![tensor("Z", &[2, 3, 4])],
+        vec![tensor("Y", &[4, 3, 2])],
+        vec![],
+        vec![
+            node(vec!["X"], vec!["Y"], "Transpose", "Transpose", vec![]),
+            node(vec!["Y"], vec!["Z"], "Transpose", "Transpose", vec![]),
+        ],
+    ));
+
+    let session =
+        pollster::block_on(wonnx::Session::from_model(model)).expect("session did not create");
+    let result = pollster::block_on(session.run(&input_data)).unwrap();
+
+    common::assert_eq_vector((&result["Z"]).try_into().unwrap(), &data);
+}
+
 #[test]
 fn test_two_transposes() {
     let mut input_data = HashMap::new();