Will need some robust batch detection. The actual batching can be handled with `vmap` but maybe slightly more efficiently with mat-mul vectorization.