From 23545e0d53e90b0965eb0a9a4960685915da772b Mon Sep 17 00:00:00 2001
From: Lyric Doshi <lyric@google.com>
Date: Sat, 7 Jan 2023 20:11:01 -0800
Subject: [PATCH] Upgrade python from 3.7 to 3.8 in the CI to address failures
 from features from newer versions.

PiperOrigin-RevId: 500460624
---
 .github/workflows/ci.yml                         | 16 ++++++++--------
 baselines/jft/input_utils.py                     | 13 +++++++------
 setup.py                                         |  2 +-
 uncertainty_baselines/datasets/toxic_comments.py |  2 +-
 4 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 7a5343f46..82b6c6f10 100644
--- a/.github/workflows/ci.yml
+++ b/.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
diff --git a/baselines/jft/input_utils.py b/baselines/jft/input_utils.py
index e34e91c58..f65d24bbd 100644
--- a/baselines/jft/input_utils.py
+++ b/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):
diff --git a/setup.py b/setup.py
index a502210f3..753569c1d 100644
--- a/setup.py
+++ b/setup.py
@@ -62,7 +62,7 @@
         ],
         'jax': [
             'clu',
-            'flax',
+            'flax==0.5.1',
             'jax',
             'jaxlib',
         ],
diff --git a/uncertainty_baselines/datasets/toxic_comments.py b/uncertainty_baselines/datasets/toxic_comments.py
index ca06e15f9..d1f6ce185 100644
--- a/uncertainty_baselines/datasets/toxic_comments.py
+++ b/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)