From e16ebe091aebcf126694c63588ccf15f8f1b3cf0 Mon Sep 17 00:00:00 2001
From: init-22 <isaac.patole022@gmail.com>
Date: Thu, 14 Nov 2024 20:33:56 +0530
Subject: [PATCH 1/8] fix: changing the dtype in random_utils to uint32

---
 algorithmic_efficiency/random_utils.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/algorithmic_efficiency/random_utils.py b/algorithmic_efficiency/random_utils.py
index cf1ea6c32..31317047e 100644
--- a/algorithmic_efficiency/random_utils.py
+++ b/algorithmic_efficiency/random_utils.py
@@ -18,8 +18,8 @@
 
 # Annoyingly, RandomState(seed) requires seed to be in [0, 2 ** 32 - 1] (an
 # unsigned int), while RandomState.randint only accepts and returns signed ints.
-MAX_INT32 = 2**31
-MIN_INT32 = -MAX_INT32
+MAX_UINT32 = 2**31
+MIN_UINT32 = 0
 
 SeedType = Union[int, list, np.ndarray]
 
@@ -35,13 +35,13 @@ def _signed_to_unsigned(seed: SeedType) -> SeedType:
 
 def _fold_in(seed: SeedType, data: Any) -> List[Union[SeedType, Any]]:
   rng = np.random.RandomState(seed=_signed_to_unsigned(seed))
-  new_seed = rng.randint(MIN_INT32, MAX_INT32, dtype=np.int32)
+  new_seed = rng.randint(MIN_UINT32, MAX_UINT32, dtype=np.uint32)
   return [new_seed, data]
 
 
 def _split(seed: SeedType, num: int = 2) -> SeedType:
   rng = np.random.RandomState(seed=_signed_to_unsigned(seed))
-  return rng.randint(MIN_INT32, MAX_INT32, dtype=np.int32, size=[num, 2])
+  return rng.randint(MIN_UINT32, MAX_UINT32, dtype=np.uint32, size=[num, 2])
 
 
 def _PRNGKey(seed: SeedType) -> SeedType:  # pylint: disable=invalid-name

From 42da4fd7bb8ee92d8fb47e4c07456ac3f2d45e3d Mon Sep 17 00:00:00 2001
From: init-22 <isaac.patole022@gmail.com>
Date: Thu, 14 Nov 2024 23:21:17 +0530
Subject: [PATCH 2/8] feat: package updates with python 3.11

---
 docker/Dockerfile | 29 +++++++++++++++-
 setup.cfg         | 86 +++++++++++++++++++++++------------------------
 2 files changed, 71 insertions(+), 44 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 9b72aea86..24d05b495 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -11,7 +11,34 @@ FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
 RUN echo "Setting up machine"
 RUN apt-get update
 RUN apt-get install -y curl tar
-RUN DEBIAN_FRONTEND=noninteractive apt-get install -y git python3 pip wget ffmpeg
+RUN DEBIAN_FRONTEND=noninteractive apt-get install -y git ffmpeg
+
+# Install prerequisites
+RUN apt-get update && apt-get install -y \
+    wget \
+    build-essential \
+    zlib1g-dev \
+    libncurses5-dev \
+    libssl-dev \
+    libreadline-dev \
+    libffi-dev \
+    curl \
+    libbz2-dev \
+    liblzma-dev
+
+# Download and install Python 3.11
+RUN cd /tmp \
+    && wget https://www.python.org/ftp/python/3.11.0/Python-3.11.0.tgz \
+    && tar -xvzf Python-3.11.0.tgz \
+    && cd Python-3.11.0 \
+    && ./configure --enable-optimizations \
+    && make -j$(nproc) \
+    && make altinstall
+
+# Create symlinks for python and pip (use 'pip' instead of 'pip3')
+RUN ln -s /usr/local/bin/python3.11 /usr/bin/python \
+    && ln -s /usr/local/bin/pip3.11 /usr/bin/pip
+
 RUN apt-get install libtcmalloc-minimal4
 RUN apt-get install unzip
 RUN apt-get install pigz
diff --git a/setup.cfg b/setup.cfg
index 4afefd164..deeb1c6c4 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -24,6 +24,7 @@ classifiers =
   Programming Language :: Python :: 3.8
   Programming Language :: Python :: 3.9
   Programming Language :: Python :: 3.10
+  Programming Language :: Python :: 3.11
   Topic :: Scientific/Engineering :: Artificial Intelligence
 
 [options]
@@ -34,22 +35,22 @@ setup_requires =
   setuptools_scm
 # Dependencies of the project:
 install_requires =
-  absl-py==1.4.0
+  absl-py==2.1.1
   # Pin to avoid unpinned install in dependencies that requires Python>=3.9.
-  networkx==3.1
-  docker==7.0.0
-  numpy>=1.23
-  pandas>=2.0.1
-  tensorflow==2.12.0
-  tensorflow-datasets==4.9.2
-  tensorflow-probability==0.20.0
-  tensorflow-addons==0.20.0
+  networkx==3.2.1
+  docker==7.1.0
+  numpy>=1.26.4
+  pandas==2.2.3
+  tensorflow==2.18.0
+  tensorflow-datasets==4.9.7
+  tensorflow-addons==0.23.0
   gputil==1.4.0
-  psutil==5.9.5
-  clu==0.0.7
-  matplotlib>=3.7.2
+  psutil==6.1.0
+  clu==0.0.12
+  matplotlib>=3.9.2
   tabulate==0.9.0
-python_requires = >=3.8
+  wandb==0.18.7
+python_requires = >=3.11
 
 
 ###############################################################################
@@ -79,78 +80,77 @@ full_dev =
 
 # Dependencies for developing the package
 dev =
-  isort==5.12.0
-  pylint==2.17.4
-  pytest==7.3.1
-  yapf==0.33.0
-  pre-commit==3.3.1
+  isort==5.13.2
+  pylint==3.3.1
+  pytest==8.3.3
+  yapf==0.43.0
+  pre-commit==4.0.1
 
 # Workloads #
 criteo1tb =
-  scikit-learn==1.2.2
+  scikit-learn==1.5.2
 
 fastmri =
-  h5py==3.8.0
-  scikit-image==0.20.0
+  h5py==3.12.1
+  scikit-image==0.24.0
 
 ogbg =
   jraph==0.0.6.dev0
-  scikit-learn==1.2.2
+  scikit-learn==1.5.2
 
 librispeech_conformer =
-  sentencepiece==0.1.99
-  tensorflow-text==2.12.1
+  sentencepiece==0.2.0
+  tensorflow-text==2.18.0
   pydub==0.25.1
 
 wmt =
-  sentencepiece==0.1.99
-  tensorflow-text==2.12.1
-  sacrebleu==1.3.1
+  sentencepiece==0.2.0
+  tensorflow-text==2.18.0
+  sacrebleu==2.4.3
 
 # Frameworks #
 
 # JAX Core
 jax_core_deps =
-  flax==0.6.10
-  optax==0.1.5
+  flax==0.10.1
+  optax==0.2.4
   # Fix chex (optax dependency) version.
   # Not fixing it can raise dependency issues with our
   # jax version.
   # Todo(kasimbeg): verify if this is necessary after we
   # upgrade jax.
-  chex==0.1.7
-  ml_dtypes==0.2.0
-  protobuf==4.25.3
+  chex==0.1.87
+  ml_dtypes==0.4.1
+  protobuf==4.25.5
 
 
 # JAX CPU
 jax_cpu =
-  jax==0.4.10
-  jaxlib==0.4.10
+  jax==0.4.35
+  jaxlib==0.4.35
   %(jax_core_deps)s
 
 # JAX GPU
 # Note this installs both jax and jaxlib.
 jax_gpu =
-  jax==0.4.10
-  jaxlib==0.4.10+cuda12.cudnn88
+  jax==0.4.35
+  jaxlib==0.4.35
+  jax-cuda12-plugin[with_cuda]==0.4.35
+  jax-cuda12-pjrt==0.4.35
   %(jax_core_deps)s
 
 # PyTorch CPU
 pytorch_cpu =
-  torch==2.1.0
-  torchvision==0.16.0
+  torch==2.5.0
+  torchvision==0.20.0
 
 # PyTorch GPU
 # Note: omit the cuda suffix and installing from the appropriate
 # wheel will result in using locally installed CUDA.
 pytorch_gpu =
-  torch==2.1.0
-  torchvision==0.16.0
+  torch==2.5.0
+  torchvision==0.20.0
 
-# wandb
-wandb =
-  wandb==0.16.5
 
 ###############################################################################
 #                           Linting Configurations                            #

From 10057769ca798f52ab6c32de41e589b45d6a5b6b Mon Sep 17 00:00:00 2001
From: init-22 <isaac.patole022@gmail.com>
Date: Thu, 14 Nov 2024 23:26:54 +0530
Subject: [PATCH 3/8] fix: reverting the python311 changes

---
 docker/Dockerfile | 29 +---------------
 setup.cfg         | 86 +++++++++++++++++++++++------------------------
 2 files changed, 44 insertions(+), 71 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 24d05b495..9b72aea86 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -11,34 +11,7 @@ FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
 RUN echo "Setting up machine"
 RUN apt-get update
 RUN apt-get install -y curl tar
-RUN DEBIAN_FRONTEND=noninteractive apt-get install -y git ffmpeg
-
-# Install prerequisites
-RUN apt-get update && apt-get install -y \
-    wget \
-    build-essential \
-    zlib1g-dev \
-    libncurses5-dev \
-    libssl-dev \
-    libreadline-dev \
-    libffi-dev \
-    curl \
-    libbz2-dev \
-    liblzma-dev
-
-# Download and install Python 3.11
-RUN cd /tmp \
-    && wget https://www.python.org/ftp/python/3.11.0/Python-3.11.0.tgz \
-    && tar -xvzf Python-3.11.0.tgz \
-    && cd Python-3.11.0 \
-    && ./configure --enable-optimizations \
-    && make -j$(nproc) \
-    && make altinstall
-
-# Create symlinks for python and pip (use 'pip' instead of 'pip3')
-RUN ln -s /usr/local/bin/python3.11 /usr/bin/python \
-    && ln -s /usr/local/bin/pip3.11 /usr/bin/pip
-
+RUN DEBIAN_FRONTEND=noninteractive apt-get install -y git python3 pip wget ffmpeg
 RUN apt-get install libtcmalloc-minimal4
 RUN apt-get install unzip
 RUN apt-get install pigz
diff --git a/setup.cfg b/setup.cfg
index deeb1c6c4..4afefd164 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -24,7 +24,6 @@ classifiers =
   Programming Language :: Python :: 3.8
   Programming Language :: Python :: 3.9
   Programming Language :: Python :: 3.10
-  Programming Language :: Python :: 3.11
   Topic :: Scientific/Engineering :: Artificial Intelligence
 
 [options]
@@ -35,22 +34,22 @@ setup_requires =
   setuptools_scm
 # Dependencies of the project:
 install_requires =
-  absl-py==2.1.1
+  absl-py==1.4.0
   # Pin to avoid unpinned install in dependencies that requires Python>=3.9.
-  networkx==3.2.1
-  docker==7.1.0
-  numpy>=1.26.4
-  pandas==2.2.3
-  tensorflow==2.18.0
-  tensorflow-datasets==4.9.7
-  tensorflow-addons==0.23.0
+  networkx==3.1
+  docker==7.0.0
+  numpy>=1.23
+  pandas>=2.0.1
+  tensorflow==2.12.0
+  tensorflow-datasets==4.9.2
+  tensorflow-probability==0.20.0
+  tensorflow-addons==0.20.0
   gputil==1.4.0
-  psutil==6.1.0
-  clu==0.0.12
-  matplotlib>=3.9.2
+  psutil==5.9.5
+  clu==0.0.7
+  matplotlib>=3.7.2
   tabulate==0.9.0
-  wandb==0.18.7
-python_requires = >=3.11
+python_requires = >=3.8
 
 
 ###############################################################################
@@ -80,77 +79,78 @@ full_dev =
 
 # Dependencies for developing the package
 dev =
-  isort==5.13.2
-  pylint==3.3.1
-  pytest==8.3.3
-  yapf==0.43.0
-  pre-commit==4.0.1
+  isort==5.12.0
+  pylint==2.17.4
+  pytest==7.3.1
+  yapf==0.33.0
+  pre-commit==3.3.1
 
 # Workloads #
 criteo1tb =
-  scikit-learn==1.5.2
+  scikit-learn==1.2.2
 
 fastmri =
-  h5py==3.12.1
-  scikit-image==0.24.0
+  h5py==3.8.0
+  scikit-image==0.20.0
 
 ogbg =
   jraph==0.0.6.dev0
-  scikit-learn==1.5.2
+  scikit-learn==1.2.2
 
 librispeech_conformer =
-  sentencepiece==0.2.0
-  tensorflow-text==2.18.0
+  sentencepiece==0.1.99
+  tensorflow-text==2.12.1
   pydub==0.25.1
 
 wmt =
-  sentencepiece==0.2.0
-  tensorflow-text==2.18.0
-  sacrebleu==2.4.3
+  sentencepiece==0.1.99
+  tensorflow-text==2.12.1
+  sacrebleu==1.3.1
 
 # Frameworks #
 
 # JAX Core
 jax_core_deps =
-  flax==0.10.1
-  optax==0.2.4
+  flax==0.6.10
+  optax==0.1.5
   # Fix chex (optax dependency) version.
   # Not fixing it can raise dependency issues with our
   # jax version.
   # Todo(kasimbeg): verify if this is necessary after we
   # upgrade jax.
-  chex==0.1.87
-  ml_dtypes==0.4.1
-  protobuf==4.25.5
+  chex==0.1.7
+  ml_dtypes==0.2.0
+  protobuf==4.25.3
 
 
 # JAX CPU
 jax_cpu =
-  jax==0.4.35
-  jaxlib==0.4.35
+  jax==0.4.10
+  jaxlib==0.4.10
   %(jax_core_deps)s
 
 # JAX GPU
 # Note this installs both jax and jaxlib.
 jax_gpu =
-  jax==0.4.35
-  jaxlib==0.4.35
-  jax-cuda12-plugin[with_cuda]==0.4.35
-  jax-cuda12-pjrt==0.4.35
+  jax==0.4.10
+  jaxlib==0.4.10+cuda12.cudnn88
   %(jax_core_deps)s
 
 # PyTorch CPU
 pytorch_cpu =
-  torch==2.5.0
-  torchvision==0.20.0
+  torch==2.1.0
+  torchvision==0.16.0
 
 # PyTorch GPU
 # Note: omit the cuda suffix and installing from the appropriate
 # wheel will result in using locally installed CUDA.
 pytorch_gpu =
-  torch==2.5.0
-  torchvision==0.20.0
+  torch==2.1.0
+  torchvision==0.16.0
 
+# wandb
+wandb =
+  wandb==0.16.5
 
 ###############################################################################
 #                           Linting Configurations                            #

From b5ad298c57659b742eb4061e9a80d152eb13abce Mon Sep 17 00:00:00 2001
From: init-22 <isaac.patole022@gmail.com>
Date: Fri, 15 Nov 2024 18:58:00 +0530
Subject: [PATCH 4/8] fix: changing PRNGkey in random_utils to key

---
 algorithmic_efficiency/random_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/algorithmic_efficiency/random_utils.py b/algorithmic_efficiency/random_utils.py
index 31317047e..91b415cd7 100644
--- a/algorithmic_efficiency/random_utils.py
+++ b/algorithmic_efficiency/random_utils.py
@@ -75,5 +75,5 @@ def split(seed: SeedType, num: int = 2) -> SeedType:
 def PRNGKey(seed: SeedType) -> SeedType:  # pylint: disable=invalid-name
   if FLAGS.framework == 'jax':
     _check_jax_install()
-    return jax_rng.PRNGKey(seed)
+    return jax_rng.key(seed)
   return _PRNGKey(seed)

From 98198682e2798f7f0a9ef6b374d6a9346ec8ec22 Mon Sep 17 00:00:00 2001
From: init-22 <isaac.patole022@gmail.com>
Date: Sat, 16 Nov 2024 12:30:47 +0530
Subject: [PATCH 5/8] fix: changing the range of MAX_UINT32

---
 algorithmic_efficiency/random_utils.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/algorithmic_efficiency/random_utils.py b/algorithmic_efficiency/random_utils.py
index 91b415cd7..93dc263bd 100644
--- a/algorithmic_efficiency/random_utils.py
+++ b/algorithmic_efficiency/random_utils.py
@@ -18,7 +18,7 @@
 
 # Annoyingly, RandomState(seed) requires seed to be in [0, 2 ** 32 - 1] (an
 # unsigned int), while RandomState.randint only accepts and returns signed ints.
-MAX_UINT32 = 2**31
+MAX_UINT32 = 2**32-1
 MIN_UINT32 = 0
 
 SeedType = Union[int, list, np.ndarray]
@@ -26,11 +26,11 @@
 
 def _signed_to_unsigned(seed: SeedType) -> SeedType:
   if isinstance(seed, int):
-    return seed % 2**32
+    return seed % MAX_UINT32
   if isinstance(seed, list):
-    return [s % 2**32 for s in seed]
+    return [s % MAX_UINT32 for s in seed]
   if isinstance(seed, np.ndarray):
-    return np.array([s % 2**32 for s in seed.tolist()])
+    return np.array([s % MAX_UINT32 for s in seed.tolist()])
 
 
 def _fold_in(seed: SeedType, data: Any) -> List[Union[SeedType, Any]]:

From 4b2e64e34e672212e4fe947674c508a410de9ef0 Mon Sep 17 00:00:00 2001
From: init-22 <isaac.patole022@gmail.com>
Date: Sat, 16 Nov 2024 12:58:26 +0530
Subject: [PATCH 6/8] bringing back PRNGKey instead of key, till the python311
 branch is merged

---
 algorithmic_efficiency/random_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/algorithmic_efficiency/random_utils.py b/algorithmic_efficiency/random_utils.py
index 93dc263bd..bcfc59c92 100644
--- a/algorithmic_efficiency/random_utils.py
+++ b/algorithmic_efficiency/random_utils.py
@@ -75,5 +75,5 @@ def split(seed: SeedType, num: int = 2) -> SeedType:
 def PRNGKey(seed: SeedType) -> SeedType:  # pylint: disable=invalid-name
   if FLAGS.framework == 'jax':
     _check_jax_install()
-    return jax_rng.key(seed)
+    return jax_rng.PRNGKey(seed)
   return _PRNGKey(seed)

From 2b8b771d169bb97a257e735ba3c26d01d6000396 Mon Sep 17 00:00:00 2001
From: init-22 <isaac.patole022@gmail.com>
Date: Tue, 19 Nov 2024 21:11:19 +0530
Subject: [PATCH 7/8] fix: ran yapf for passing the checks

---
 algorithmic_efficiency/random_utils.py | 2 +-
 scoring/score_submissions.py           | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/algorithmic_efficiency/random_utils.py b/algorithmic_efficiency/random_utils.py
index bcfc59c92..f40a98003 100644
--- a/algorithmic_efficiency/random_utils.py
+++ b/algorithmic_efficiency/random_utils.py
@@ -18,7 +18,7 @@
 
 # Annoyingly, RandomState(seed) requires seed to be in [0, 2 ** 32 - 1] (an
 # unsigned int), while RandomState.randint only accepts and returns signed ints.
-MAX_UINT32 = 2**32-1
+MAX_UINT32 = 2**32 - 1
 MIN_UINT32 = 0
 
 SeedType = Union[int, list, np.ndarray]
diff --git a/scoring/score_submissions.py b/scoring/score_submissions.py
index 59295b686..8cc06b15f 100644
--- a/scoring/score_submissions.py
+++ b/scoring/score_submissions.py
@@ -211,7 +211,8 @@ def main(_):
         verbosity=0,
         self_tuning_ruleset=FLAGS.self_tuning_ruleset,
         strict=FLAGS.strict,
-        output_dir=FLAGS.output_dir,)
+        output_dir=FLAGS.output_dir,
+    )
     if not os.path.exists(FLAGS.output_dir):
       os.mkdir(FLAGS.output_dir)
     performance_profile.plot_performance_profiles(

From ff176d7719d55238a898f29c3fe5e3812bd6a06c Mon Sep 17 00:00:00 2001
From: init-22 <isaac.patole022@gmail.com>
Date: Fri, 22 Nov 2024 22:28:00 +0530
Subject: [PATCH 8/8] fix: triggering the checks again

---
 scoring/performance_profile.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/scoring/performance_profile.py b/scoring/performance_profile.py
index 0d5ca9770..f4f2d5679 100644
--- a/scoring/performance_profile.py
+++ b/scoring/performance_profile.py
@@ -321,7 +321,7 @@ def compute_performance_profiles(submissions,
   df = df[BASE_WORKLOADS + HELDOUT_WORKLOADS]
   # Sort workloads alphabetically (for better display)
   df = df.reindex(sorted(df.columns), axis=1)
-  
+
   # Save time to target dataframe
   df.to_csv(os.path.join(output_dir, 'time_to_targets.csv'))
   # For each held-out workload set to inf if the base workload is inf or nan