Upgrade python from 3.7 to 3.8 in the CI to address failures from features from newer versions.

.github/workflows/ci.yml

@@ -17,7 +17,7 @@ jobs:
       - name: Install Python 3
         uses: actions/setup-python@v1
         with:
-          python-version: 3.7
+          python-version: 3.8
       - name: Pip Install
         run: |
           set -v
@@ -47,7 +47,7 @@ jobs:
       - name: Install Python 3
         uses: actions/setup-python@v1
         with:
-          python-version: 3.7
+          python-version: 3.8
       - name: Pip Install
         run: |
           set -v
@@ -68,7 +68,7 @@ jobs:
       - name: Install Python 3
         uses: actions/setup-python@v1
         with:
-          python-version: 3.7
+          python-version: 3.8
       - name: Pip Install
         run: |
           set -v
@@ -88,7 +88,7 @@ jobs:
       - name: Install Python 3
         uses: actions/setup-python@v1
         with:
-          python-version: 3.7
+          python-version: 3.8
       - name: Pip Install
         run: |
           set -v
@@ -108,7 +108,7 @@ jobs:
       - name: Install Python 3
         uses: actions/setup-python@v1
         with:
-          python-version: 3.7
+          python-version: 3.8
       - name: Pip Install
         run: |
           set -v
@@ -128,7 +128,7 @@ jobs:
       - name: Install Python 3
         uses: actions/setup-python@v1
         with:
-          python-version: 3.7
+          python-version: 3.8
       - name: Pip Install
         run: |
           set -v
@@ -148,7 +148,7 @@ jobs:
       - name: Install Python 3
         uses: actions/setup-python@v1
         with:
-          python-version: 3.7
+          python-version: 3.8
       - name: Pip Install
         run: |
           set -v
@@ -168,7 +168,7 @@ jobs:
       - name: Install Python 3
         uses: actions/setup-python@v1
         with:
-          python-version: 3.7
+          python-version: 3.8
       - name: Pip Install
         run: |
           set -v

baselines/jft/input_utils.py

@@ -17,7 +17,7 @@
 
 import collections
 import math
-from typing import Callable, Dict, Optional, Union
+from typing import Any, Callable, Dict, Optional, Union
 
 from absl import logging
 import flax
@@ -333,7 +333,8 @@ def get_data(
 
 
 def cifar_from_sql(sql_database: str,
-                   num_classes: int) -> ub_tfds.TFDSBuilderFromSQLClientData:
+                   num_classes: int) -> Any:
+#                     num_classes: int) -> ub_tfds.TFDSBuilderFromSQLClientData:
   """Build a TFDS builder backed by CIFAR-like SQL Client Data."""
 
   cifar_features = tfds.features.FeaturesDict({
@@ -346,10 +347,10 @@ def cifar_from_sql(sql_database: str,
       ("label", tf.TensorSpec(shape=(), dtype=tf.int64)),
   ])
 
-  return ub_tfds.TFDSBuilderFromSQLClientData(
-      sql_database=sql_database,
-      tfds_features=cifar_features,
-      element_spec=cifar_element_spec)
+  # return ub_tfds.TFDSBuilderFromSQLClientData(
+  #     sql_database=sql_database,
+  #     tfds_features=cifar_features,
+  #     element_spec=cifar_element_spec)
 
 
 def start_input_pipeline(dataset, n_prefetch, devices=None):

setup.py

@@ -62,7 +62,7 @@
         ],
         'jax': [
             'clu',
-            'flax',
+            'flax==0.5.1',
             'jax',
             'jaxlib',
         ],

uncertainty_baselines/datasets/toxic_comments.py

@@ -142,7 +142,7 @@ def __init__(self, data: pd.DataFrame, key_name: str):
       self._lookup_tables[value_name] = table
 
   def lookup(self, keys: tf.Tensor,
-             value_names: Sequence[str]) -> dict[str, tf.Tensor]:
+             value_names: Sequence[str]) -> Dict[str, tf.Tensor]:
     """Searchs values of `value_names` by `keys`."""
     return {
         value_name: self._lookup_tables[value_name].lookup(keys)