()
| 1221 | test_trainium_layout_0() |
| 1222 | |
| 1223 | def test_trainium_layout_1(): |
| 1224 | layout = TileLayout(S[(2, 6, 4, 2, 2) : (1 @ F, 1 @ P, 12 @ F, 6 @ P, 48 @ F)]) |
| 1225 | |
| 1226 | def f(i0, i1): |
| 1227 | leaf1 = i0 // 6 |
| 1228 | leaf2 = i0 % 6 |
| 1229 | leaf3 = i1 // 4 |
| 1230 | leaf4 = (i1 % 4) // 2 |
| 1231 | leaf5 = i1 % 2 |
| 1232 | coord = layout.apply(i0, i1, shape=(12, 16)) |
| 1233 | assert coord["P"] == leaf2 + leaf4 * 6 |
| 1234 | assert coord["F"] == leaf1 * 1 + leaf3 * 12 + leaf5 * 48 |
| 1235 | |
| 1236 | for i0, i1 in itertools.product(range(6), range(16)): |
| 1237 | f(i0, i1) |
| 1238 | for i in range(6 * 16): |
| 1239 | f(i // 16, i % 16) |
| 1240 | |
| 1241 | test_trainium_layout_1() |
| 1242 |
no test coverage detected
searching dependent graphs…