All kinds of strange things when batch size is not equal to max_batch_size

[dfdx]

JAX is quite strict about buffer sizes, and so it's generally expected that a user applies models with batch size equal to ModelArgs.max_batch_size. Yet, we neither forbid different batch sizes, nor handle it properly. In some cases, different batch size leads to matrix size issues, in other cases the returned array has incorrect size, etc.

We need to either validate input tensors, or handle cases of batch size < max_batch_size properly.