[SYSTEMDS-3803] DML-bodied Util Function for Transposing ABCD to ACBD

This patch adds a simple util function for transposing matrices in a
specified way, which is required for multi-head attention implementation.

Closes #2151

Author: Maximilian.S

scripts/nn/util.dml

@@ -380,3 +380,28 @@ top_k2d = function(matrix[double] X, int k, int C, int Hin, int Win)
   indices = transpose_NCHW_to_CNHW(indices_K_NHW, N)
 }
 
+transpose_ABCD_to_ACBD = function(matrix[double] X, int B, int C)
+    return (matrix[double] out) {
+  /*
+   * Reshape util for tensors in ABCD format.
+   * Transposes the 2nd and 3rd axes.
+   *
+   * Inputs:
+   *  - X: Inputs, of shape (A, B*C*D).
+   *  - B: Dimension of 2nd axis.
+   *  - C: Dimension of 3rd axis.
+   *
+   * Outputs:
+   *  - out: Outputs with the 2nd and 3rd axes transposed, of
+   *      shape (A, C*B*D).
+   */
+  A = nrow(X)
+  BCD = ncol(X)
+
+  # use NCHW_to_CNHW for X: (A, B*C*D) -> (B, A*C*D)
+  X_BACD = transpose_NCHW_to_CNHW(X, B)
+  # use NCHW_to_CNHW for X: (B, A*C*D) -> (A*C, B*D)
+  X_ACBD = transpose_NCHW_to_CNHW(X_BACD, A*C)
+  # reshape X: (A*C, B*D) -> (A, C*B*D)
+  out = matrix(X_ACBD, rows=A, cols=BCD)
+}

src/test/java/org/apache/sysds/test/applications/nn/NNComponentTest.java

@@ -108,6 +108,11 @@ public void transpose_NCHW_to_CNHW() {
 		run("transpose_NCHW_to_CNHW.dml");
 	}
 
+	@Test
+	public void transpose_ABCD_to_ACBD() {
+		run("transpose_ABCD_to_ACBD.dml");
+	}
+
 	@Test 
 	public void logcosh(){
 		run("logcosh.dml");

src/test/scripts/applications/nn/component/transpose_ABCD_to_ACBD.dml

@@ -0,0 +1,82 @@
+#-------------------------------------------------------------
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+#-------------------------------------------------------------
+
+source("src/test/scripts/applications/nn/util.dml") as test_util
+source("scripts/nn/util.dml") as util
+
+
+transpose_ABCD_to_ACBD = function() {
+  /*
+   * Test for `transpose_ABCD_to_ACBD` function.
+   */
+  print("Testing transpose_ABCD_to_ACBD function.")
+
+  # Generate data
+  A = 2
+  B = 3
+  C = 4
+  D = 5
+  X = matrix(seq(1, A*B*C*D), rows=A, cols=B*C*D)
+
+  out = util::transpose_ABCD_to_ACBD(X, B, C)
+
+  target =
+    matrix("1 2 3 4 5 21 22 23 24 25 41 42 43 44 45
+            6 7 8 9 10 26 27 28 29 30 46 47 48 49 50
+            11 12 13 14 15 31 32 33 34 35 51 52 53 54 55
+            16 17 18 19 20 36 37 38 39 40 56 57 58 59 60
+
+            61 62 63 64 65 81 82 83 84 85 101 102 103 104 105
+            66 67 68 69 70 86 87 88 89 90 106 107 108 109 110
+            71 72 73 74 75 91 92 93 94 95 111 112 113 114 115
+            76 77 78 79 80 96 97 98 99 100 116 117 118 119 120",
+           rows=A, cols=C*B*D)
+
+  # Equivalency check
+  test_util::check_all_close(out, target, 1e-10)
+}
+
+
+transpose_ABCD_to_ACBD_single_val = function() {
+  /*
+   * Test for `transpose_ABCD_to_ACBD` function,
+   * transposing a single value matrix.
+   */
+  print("Testing transpose_ABCD_to_ACBD function with single value.")
+
+  # Generate data
+  A = 1
+  B = 1
+  C = 1
+  D = 1
+  X = matrix(seq(1, A*B*C*D), rows=A, cols=B*C*D)
+
+  out = util::transpose_ABCD_to_ACBD(X, B, C)
+
+  target = X
+
+  # Equivalency check
+  test_util::check_all_close(out, target, 1e-10)
+}
+
+
+transpose_ABCD_to_ACBD()
+transpose_ABCD_to_ACBD_single_val()