| 118 | } |
| 119 | |
| 120 | static array dist_tile2(array a, array b) { |
| 121 | int feat_len = (int)a.dims(0); |
| 122 | int alen = (int)a.dims(1); |
| 123 | int blen = (int)b.dims(1); |
| 124 | |
| 125 | // Shape of a is (feat_len, alen, 1) |
| 126 | array a_mod = a; |
| 127 | // Reshape b from (feat_len, blen) to (feat_len, 1, blen) |
| 128 | array b_mod = moddims(b, feat_len, 1, blen); |
| 129 | |
| 130 | // Tile both matrices to be (feat_len, alen, blen) |
| 131 | array a_tiled = tile(a_mod, 1, 1, blen); |
| 132 | array b_tiled = tile(b_mod, 1, alen, 1); |
| 133 | |
| 134 | // Do The sum operation along first dimension |
| 135 | // Output is of shape (1, alen, blen) |
| 136 | array dist_mod = sum(abs(a_tiled - b_tiled)); |
| 137 | |
| 138 | // Reshape dist_mat from (1, alen, blen) to (alen, blen) |
| 139 | array dist_mat = moddims(dist_mod, alen, blen); |
| 140 | return dist_mat; |
| 141 | } |
| 142 | |
| 143 | static void bench_naive() { dist_naive(A, B); } |
| 144 | |