Batch Matmul

Даны два батча матриц (3D списки):
- A: размер (B, N, K)
- B_mat: размер (B, K, M)

Вычислите поэлементное (по батчу) матричное произведение.

Сигнатура

def batch_matmul(a: list[list[list[int]]], b: list[list[list[int]]]) -> list[list[list[int]]]:

Примеры

a = [[[1,2],[3,4]], [[5,6],[7,8]]]
b = [[[1,0],[0,1]], [[2,0],[0,2]]]
batch_matmul(a, b) → [[[1,2],[3,4]], [[10,12],[14,16]]]