vit -> 27 errors

Author: enpasos

jaxamples/mnist_config.py

@@ -309,10 +309,10 @@ def shared_mnist_augmentation_config() -> AugmentationConfig:
         scale_max_x=1.1,
         scale_min_y=0.9,
         scale_max_y=1.1,
-        scaling_probability=0.5,
+        scaling_probability=0.6,
         enable_rotation=True,
         max_rotation=12.0,
-        rotation_probability=0.5,
+        rotation_probability=0.7,
         enable_elastic=True,
         elastic_alpha=1.2,
         elastic_sigma=1.0,

onnx/mnist_vit_model.onnx

@@ -58,12 +58,12 @@
       "rect_erase_height": 2,
       "rect_erase_width": 20,
       "rect_erasing_probability": 0.0,
-      "rotation_probability": 0.5,
+      "rotation_probability": 0.7,
       "scale_max_x": 1.1,
       "scale_max_y": 1.1,
       "scale_min_x": 0.9,
       "scale_min_y": 0.9,
-      "scaling_probability": 0.5,
+      "scaling_probability": 0.6,
       "translation_probability": 0.8
     },
     "base_learning_rate": 0.0001,

onnx/mnist_vit_model_config.json

@@ -47,12 +47,12 @@ def test_mnist_cnn_default_config_uses_shared_augmentation():
     assert cnn_config.training.augmentation.scale_max_x == pytest.approx(1.1)
     assert cnn_config.training.augmentation.scale_min_y == pytest.approx(0.9)
     assert cnn_config.training.augmentation.scale_max_y == pytest.approx(1.1)
-    assert cnn_config.training.augmentation.scaling_probability == pytest.approx(0.5)
+    assert cnn_config.training.augmentation.scaling_probability == pytest.approx(0.6)
     assert cnn_config.training.augmentation.max_rotation == pytest.approx(12.0)
-    assert cnn_config.training.augmentation.rotation_probability == pytest.approx(0.5)
-    assert cnn_config.training.augmentation.elastic_alpha == pytest.approx(1.0)
+    assert cnn_config.training.augmentation.rotation_probability == pytest.approx(0.7)
+    assert cnn_config.training.augmentation.elastic_alpha == pytest.approx(1.2)
     assert cnn_config.training.augmentation.elastic_sigma == pytest.approx(1.0)
-    assert cnn_config.training.augmentation.elastic_probability == pytest.approx(0.2)
+    assert cnn_config.training.augmentation.elastic_probability == pytest.approx(0.3)
     assert cnn_config.training.checkpoint_dir.endswith("cnn_c32_64_128_128_d256_checkpoints")
     assert cnn_config.model.conv_channels == [32, 64, 128, 128]
     assert cnn_config.model.dense_hidden_dim == 256