From ea4e284e9c7b0e5d59ea0a9cf1b9d5c67a0d59be Mon Sep 17 00:00:00 2001
From: mark <markblee@apple.com>
Date: Sat, 16 Mar 2024 17:20:56 -0700
Subject: [PATCH] Bump tensorflow to 2.16.1 and tensorstore to >= 0.1.56.

---
 axlearn/cloud/gcp/vertexai_tensorboard.py |  3 +--
 axlearn/common/quantizer_test.py          |  2 +-
 axlearn/vision/augment.py                 |  4 ++--
 docs/01-start.md                          | 11 +++++++++--
 pyproject.toml                            |  8 ++++----
 5 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/axlearn/cloud/gcp/vertexai_tensorboard.py b/axlearn/cloud/gcp/vertexai_tensorboard.py
index b3e6a0bc..99441bfa 100644
--- a/axlearn/cloud/gcp/vertexai_tensorboard.py
+++ b/axlearn/cloud/gcp/vertexai_tensorboard.py
@@ -11,7 +11,6 @@
     from google.cloud.aiplatform import initializer
     from google.cloud.aiplatform.tensorboard import uploader, uploader_utils
     from google.cloud.aiplatform.utils import TensorboardClientWithOverride
-    from tensorboard.uploader import util
 
     _VERTEXAI_INSTALLED = True
 except (ImportError, ModuleNotFoundError):
@@ -74,7 +73,7 @@ def _start_vertexai_tensorboard(*, project_id: str, region: str, resource_name:
         event_file_inactive_secs=None,
         verbosity=0,
         run_name_prefix=None,
-        logdir_poll_rate_limiter=util.RateLimiter(interval_secs=30),
+        logdir_poll_rate_limiter=uploader_utils.RateLimiter(interval_secs=30),
     )
     tb_uploader.create_experiment()
     tb_uploader.start_uploading()
diff --git a/axlearn/common/quantizer_test.py b/axlearn/common/quantizer_test.py
index fde03b53..7707aeb3 100644
--- a/axlearn/common/quantizer_test.py
+++ b/axlearn/common/quantizer_test.py
@@ -7,6 +7,7 @@
 # Licensed under the Apache License, Version 2.0 (the "License").
 
 """Tests quantization layers and metrics."""
+
 # pylint: disable=no-self-use,wrong-import-position,missing-module-docstring
 from typing import List
 
@@ -211,7 +212,6 @@ def test_forward(
             ),
             shapes(layer_params),
         )
-        # pylint: disable-next=protected-access
         with prng_impl("threefry2x32"):
             proj_key = _create_prngkeyarray([3077990774, 2166202870])
             codebook_key = _create_prngkeyarray([791337683, 1373966058])
diff --git a/axlearn/vision/augment.py b/axlearn/vision/augment.py
index 7e5f775b..6a84df9f 100644
--- a/axlearn/vision/augment.py
+++ b/axlearn/vision/augment.py
@@ -19,12 +19,12 @@
 Reference:
 https://github.com/tensorflow/models/blob/master/official/legacy/image_classification/augment.py
 """
+
 # pylint: disable=too-many-lines
 import math
 from typing import Any, List, Optional, Sequence, Text, Tuple, Union
 
 import tensorflow as tf
-from keras import backend
 
 # This signifies the max integer that the controller RNN could predict for the
 # augmentation scheme.
@@ -178,7 +178,7 @@ def _keras_image_processing_transform(
     https://github.com/keras-team/keras/blob/v2.14.0/keras/layers/preprocessing/image_preprocessing.py#L720-L815
     since it was removed from the public API, with minor adaptation.
     """
-    with backend.name_scope(name or "transform"):
+    with tf.name_scope(name or "transform"):
         if output_shape is None:
             output_shape = tf.shape(images)[1:3]
             if not tf.executing_eagerly():
diff --git a/docs/01-start.md b/docs/01-start.md
index 9441227a..5084eac8 100644
--- a/docs/01-start.md
+++ b/docs/01-start.md
@@ -48,9 +48,16 @@ conda install -c apple tensorflow-deps
 # Manually build tensorflow-text until a collaborator build is available.
 # This was tested using clang version 15 - you may get non-working wheels with earlier versions of clang.
 mkdir ~/builds && git clone https://github.com/tensorflow/text.git ~/builds/text
-cd ~/builds/text && git checkout 6064f1bf8fd078777b6c8690986b908c28764a94
+# Install tensorflow prior to building.
+pip install tensorflow==2.16.1
+cd ~/builds/text && git checkout 0f9f6df5b4da19bc7a734ba05fc4fa12bccbedbe
+
+# If you do NOT have bazel installed, run the following and ensure that $HOME/bin is in $PATH:
+# ./oss_scripts/install_bazel.sh
+
+# Build tensorflow-text.
 ./oss_scripts/run_build.sh
-pip install ./tensorflow_text-2.14.0-cp39-cp39-macosx_*_arm64.whl
+pip install ./tensorflow_text-2.16.1-cp39-cp39-macosx_*_arm64.whl
 ```
 </details>
 
diff --git a/pyproject.toml b/pyproject.toml
index 8145bbb7..16e6820b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -23,14 +23,14 @@ dependencies = [
     "optax==0.1.7",  # optimizers (0.1.0 has known bugs).
     "portpicker",
     "protobuf>=3.20.3",
-    "tensorboard-plugin-profile==2.14.0",
+    "tensorboard-plugin-profile==2.15.1",
     # This has both x86 and arm64 wheels. Underneath the hood it uses tensorflow-macos since 2.13.
-    "tensorflow==2.14.0",
+    "tensorflow==2.16.1",
     "tensorflow-datasets>=4.9.2",
     # tensorflow-io 0.34.0 enabled macOS arm64 builds.
     "tensorflow-io>=0.34.0",
-    "tensorflow_text==2.14.0; platform_machine == 'x86_64'", # implied by seqio, but also used directly for text processing
-    "tensorstore>=0.1.21",  # used for supporting GDA checkpoints
+    "tensorflow_text==2.16.1; platform_machine == 'x86_64'", # implied by seqio, but also used directly for text processing
+    "tensorstore>=0.1.56",  # used for supporting GDA checkpoints
     "toml",  # for config management
     "typing-extensions==4.9.0",  # needed for typing.Protocol. >4.9.0 runs into attribute error `__non_callable_proto_members__`.
     "scipy==1.12.0",  # to avoid "module 'scipy.linalg' has no attribute 'tril'"