better tests for weight loading with refactoring (#219)

* move better tests from PR 214 to investigate weight folding issues

[PR 214 with tokenizer updates](#214)
appears to have exposed a bug. I made some better tests on that branch,
and I'm adding them to the main branch to see if that also breaks.

* format

* update deps

* indexed tokenizer tests for save/load

maze_transformer/test_helpers/assertions.py

@@ -93,7 +93,8 @@ def assert_model_output_equality(
         )
         if not torch.all(output_argsort_match):
             raise ModelOutputArgsortEqualityError(
-                f"argsort not equal, {output_argsort_match.numel() - output_argsort_match.sum()} / {output_argsort_match.numel()} elements differ"
+                f"argsort not equal, {output_argsort_match.numel() - output_argsort_match.sum()} / {output_argsort_match.numel()} elements differ",
+                f"{vocab_size = }, {test_sequence_length = }",
             )
 
     # apply normalization (e.g. softmax) and check with atol v-small