| 359 | |
| 360 | |
| 361 | def matmul4d(x1, x2): |
| 362 | batchs, heads = x1.shape[0], x1.shape[1] |
| 363 | ys = [] |
| 364 | for b in range(batchs): |
| 365 | x1b, x2b = autograd.squeeze(x1[b]), autograd.squeeze(x2[b]) |
| 366 | yb = [] |
| 367 | for h in range(heads): |
| 368 | x1h, x2h = autograd.squeeze(x1b[h]), autograd.squeeze(x2b[h]) |
| 369 | yh = autograd.matmul(x1h, x2h) |
| 370 | yh = autograd.unsqueeze(yh, axis=[0]) |
| 371 | yb.append(yh) |
| 372 | yb = autograd.cat(yb, axis=0) |
| 373 | yb = autograd.unsqueeze(yb, axis=[0]) |
| 374 | ys.append(yb) |
| 375 | y = autograd.cat(ys, axis=0) |
| 376 | return y |
| 377 | |
| 378 | |
| 379 | class MultiHeadAttention(layer.Layer): |