Fix AQT + Jax 0.4.16 by removing preferred_element_type assertion (#189)

* Remove preferred_element_type assertion to fix AQT on jax 0.4.16

.github/workflows/UnitTests.yml

@@ -56,3 +56,7 @@ jobs:
       run: |
         source venv/bin/activate
         python3 MaxText/decode.py MaxText/configs/base.yml run_name=runner_$(date +%Y-%m-%d-%H-%M) base_output_directory=gs://runner-maxtext-logs dataset_path=gs://maxtext-dataset steps=2
+    - name: Test int8_training
+      run: |
+        source venv/bin/activate
+        python3 MaxText/train.py MaxText/configs/base.yml run_name=runner_$(date +%Y-%m-%d-%H-%M) base_output_directory=gs://runner-maxtext-logs dataset_path=gs://maxtext-dataset int8_training=true steps=2

MaxText/aqt/jax/v2/aqt_dot_general.py

@@ -540,10 +540,10 @@ def ret_dg(
     assert (
         precision is None
     ), f'Precision {precision} requested together with quantization.'
-    assert preferred_element_type is None, (
-        f'Preferred_element_typerecision {preferred_element_type} requested'
-        ' together with quantization.'
-    )
+
+    # The quantized einsum ignores preferred_element_type (b/302728979)
+    preferred_element_type = None
+
     assert lhs.dtype == rhs.dtype, (
         'The only reason we need that, is because we need to determine return'
         ' type.'