Feature/larger batches

Previously, the number of batches was limited to 128. This PR changes the `matrix_offsets` logic to fix it.