Implement working with batches as a decorator

[Bihaqo]

Something like

@support_batches
def matmul(a, b):
  ...
matmul(a, b)
matmul(batch_a, b)
matmul(batch_a, batch_b)