Temp testing automation, please ignore

.github/workflows/pypi-publish.yml

@@ -2,17 +2,21 @@ name: pypi
 
 on:
   release:
-    types: [created]
-
+    types: [created, published]
+    branches: [main, master]
 jobs:
   deploy:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v4
     - name: Set up Python
-      uses: actions/setup-python@v1
+      uses: actions/setup-python@v4
       with:
         python-version: '3.x'
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install setuptools wheel twine
     - name: Check consistency between the package version and release tag
       run: |
         RELEASE_VER=${GITHUB_REF#refs/*/}
@@ -21,10 +25,6 @@ jobs:
         then
           echo "package ver. ($PACKAGE_VER) != release ver. ($RELEASE_VER)"; exit 1
         fi
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        pip install setuptools wheel twine
     - name: Build and publish
       env:
         TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}

README.md

@@ -20,7 +20,7 @@ pip install dm-clrs
 or directly from GitHub (updated more frequently):
 
 ```shell
-pip install git+git://github.com/deepmind/clrs.git
+pip install git+https://github.com/google-deepmind/clrs.git
 ```
 
 You may prefer to install it in a virtual environment if any requirements
@@ -29,7 +29,7 @@ clash with your Python installation:
 ```shell
 python3 -m venv clrs_env
 source clrs_env/bin/activate
-pip install git+git://github.com/deepmind/clrs.git
+pip install git+https://github.com/google-deepmind/clrs.git
 ```
 
 Once installed you can run our example baseline model:
@@ -225,6 +225,14 @@ for feedback in _iterate_sampler(batch_size=32):
 
 ```
 
+Most recently, we are offering [**CLRS-Text**](https://github.com/google-deepmind/clrs/tree/master/clrs/_src/clrs_text),
+a text-based variant of the benchmark suitable for training and evaluating the algorithmic reasoning
+capabilities of language models. Please see the relevant subfolder for a
+dedicated README file.
+
+You may also see the [companion paper](https://arxiv.org/abs/2406.04229) on
+CLRS-Text.
+
 ## Adding new algorithms
 
 Adding a new algorithm to the task suite requires the following steps:
@@ -259,3 +267,16 @@ To cite the CLRS Algorithmic Reasoning Benchmark:
   year={2022}
 }
 ```
+
+To cite the CLRS-Text Algorithmic Reasoning Language Benchmark:
+
+```latex
+@article{deepmind2024clrstext,
+  title={The CLRS-Text Algorithmic Reasoning Language Benchmark},
+  author={Larisa Markeeva and Sean McLeish and Borja Ibarz and Wilfried Bounsi
+    and Olga Kozlova and Alex Vitvitskyi and Charles Blundell and
+    Tom Goldstein and Avi Schwarzschild and Petar Veli\v{c}kovi\'{c}},
+  journal={arXiv preprint arXiv:2406.04229},
+  year={2024}
+}
+```

clrs/__init__.py

@@ -16,21 +16,30 @@
 """The CLRS Algorithmic Reasoning Benchmark."""
 
 from clrs import models
+
 from clrs._src import algorithms
+from clrs._src import clrs_text
 from clrs._src import decoders
 from clrs._src import processors
+from clrs._src import specs
+
 from clrs._src.dataset import chunkify
 from clrs._src.dataset import CLRSDataset
 from clrs._src.dataset import create_chunked_dataset
 from clrs._src.dataset import create_dataset
 from clrs._src.dataset import get_clrs_folder
 from clrs._src.dataset import get_dataset_gcp_url
+
 from clrs._src.evaluation import evaluate
 from clrs._src.evaluation import evaluate_hints
+
 from clrs._src.model import Model
+
 from clrs._src.probing import DataPoint
 from clrs._src.probing import predecessor_to_cyclic_predecessor_and_first
+
 from clrs._src.processors import get_processor_factory
+
 from clrs._src.samplers import build_sampler
 from clrs._src.samplers import CLRS30
 from clrs._src.samplers import Features
@@ -40,6 +49,7 @@
 from clrs._src.samplers import process_random_pos
 from clrs._src.samplers import Sampler
 from clrs._src.samplers import Trajectory
+
 from clrs._src.specs import ALGO_IDX_INPUT_NAME
 from clrs._src.specs import CLRS_30_ALGS_SETTINGS
 from clrs._src.specs import Location
@@ -49,7 +59,7 @@
 from clrs._src.specs import Stage
 from clrs._src.specs import Type
 
-__version__ = "1.0.0"
+__version__ = "2.0.1"
 
 __all__ = (
     "ALGO_IDX_INPUT_NAME",
@@ -59,6 +69,7 @@
     "CLRS_30_ALGS_SETTINGS",
     "create_chunked_dataset",
     "create_dataset",
+    "clrs_text",
     "get_clrs_folder",
     "get_dataset_gcp_url",
     "get_processor_factory",
@@ -67,6 +78,7 @@
     "process_permutations",
     "process_pred_as_input",
     "process_random_pos",
+    "specs",
     "evaluate",
     "evaluate_hints",
     "Features",

clrs/_src/algorithms/searching.py

@@ -179,8 +179,10 @@ def partition(A, A_pos, p, r, target, probes):
               'i': probing.mask_one(A_pos[i + 1], A.shape[0]),
               'j': probing.mask_one(A_pos[j], A.shape[0]),
               'i_rank': (i + 1) * 1.0 / A.shape[0],
-              'target': target * 1.0 / A.shape[0]
-          })
+              'target': target * 1.0 / A.shape[0],
+              'pivot': probing.mask_one(A_pos[r], A.shape[0]),
+          },
+      )
 
     tmp = A[i + 1]
     A[i + 1] = A[r]
@@ -199,8 +201,10 @@ def partition(A, A_pos, p, r, target, probes):
             'i': probing.mask_one(A_pos[i + 1], A.shape[0]),
             'j': probing.mask_one(A_pos[r], A.shape[0]),
             'i_rank': (i + 1 - p) * 1.0 / A.shape[0],
-            'target': target * 1.0 / A.shape[0]
-        })
+            'target': target * 1.0 / A.shape[0],
+            'pivot': probing.mask_one(A_pos[i + 1], A.shape[0]),
+        },
+    )
 
     return i + 1
 

clrs/_src/algorithms/sorting.py

@@ -123,13 +123,8 @@ def bubble_sort(A: _Array) -> _Out:
   for i in range(A.shape[0] - 1):
     for j in reversed(range(i + 1, A.shape[0])):
       if A[j] < A[j - 1]:
-        tmp = A[j]
-        A[j] = A[j - 1]
-        A[j - 1] = tmp
-
-        tmp = A_pos[j]
-        A_pos[j] = A_pos[j - 1]
-        A_pos[j - 1] = tmp
+        A[j], A[j - 1] = A[j - 1], A[j]
+        A_pos[j], A_pos[j - 1] = A_pos[j - 1], A_pos[j]
 
       probing.push(
           probes,
@@ -190,13 +185,8 @@ def max_heapify(A, i, heap_size, ind, phase):
     if r < heap_size and A[r] > A[largest]:
       largest = r
     if largest != i:
-      tmp = A[i]
-      A[i] = A[largest]
-      A[largest] = tmp
-
-      tmp = A_pos[i]
-      A_pos[i] = A_pos[largest]
-      A_pos[largest] = tmp
+      A[i], A[largest] = A[largest], A[i]
+      A_pos[i], A_pos[largest] = A_pos[largest], A_pos[i]
 
     probing.push(
         probes,
@@ -221,13 +211,8 @@ def build_max_heap(A):
   build_max_heap(A)
   heap_size = A.shape[0]
   for i in reversed(range(1, A.shape[0])):
-    tmp = A[0]
-    A[0] = A[i]
-    A[i] = tmp
-
-    tmp = A_pos[0]
-    A_pos[0] = A_pos[i]
-    A_pos[i] = tmp
+    A[0], A[i] = A[i], A[0]
+    A_pos[0], A_pos[i] = A_pos[i], A_pos[0]
 
     heap_size -= 1
 
@@ -268,12 +253,8 @@ def partition(A, A_pos, p, r, probes):
     for j in range(p, r):
       if A[j] <= x:
         i += 1
-        tmp = A[i]
-        A[i] = A[j]
-        A[j] = tmp
-        tmp = A_pos[i]
-        A_pos[i] = A_pos[j]
-        A_pos[j] = tmp
+        A[i], A[j] = A[j], A[i]
+        A_pos[i], A_pos[j] = A_pos[j], A_pos[i]
 
       probing.push(
           probes,
@@ -286,12 +267,8 @@ def partition(A, A_pos, p, r, probes):
               'j': probing.mask_one(A_pos[j], A.shape[0])
           })
 
-    tmp = A[i + 1]
-    A[i + 1] = A[r]
-    A[r] = tmp
-    tmp = A_pos[i + 1]
-    A_pos[i + 1] = A_pos[r]
-    A_pos[r] = tmp
+    A[i + 1], A[r] = A[r], A[i + 1]
+    A_pos[i + 1], A_pos[r] = A_pos[r], A_pos[i + 1]
 
     probing.push(
         probes,

clrs/_src/baselines.py

@@ -766,7 +766,9 @@ def _keep_in_algo(k, v):
     masked_grads = grads
   else:
     masked_grads = {k: _keep_in_algo(k, v) for k, v in grads.items()}
-  flat_grads, treedef = jax.tree_util.tree_flatten(masked_grads)
+  flat_grads, treedef = jax.tree_util.tree_flatten(
+      masked_grads, is_leaf=lambda x: x is None
+  )
   flat_opt_state = jax.tree_util.tree_map(
       lambda _, x: x  # pylint:disable=g-long-lambda
       if isinstance(x, (np.ndarray, jax.Array))

clrs/_src/clrs_text/README.md

@@ -0,0 +1,57 @@
+# The CLRS-Text Algorithmic Reasoning Language Benchmark
+
+Eliciting reasoning capabilities from language models (LMs) is a critical
+direction on the path towards building intelligent systems. Most recent studies dedicated to reasoning focus on out-of-distribution performance on
+procedurally-generated synthetic benchmarks, bespoke-built to evaluate specific
+skills only. One standard way in which OOD performance is measured, for example,
+is by assessing [length generalisation](https://arxiv.org/abs/2402.09371).
+
+CLRS-Text is a textual version of the traces generated by thirty
+algorithms selected from the third edition of the standard
+*Introduction to Algorithms* textbook by Cormen, Leiserson, Rivest and Stein. It
+serves to consolidate and unify previous lines of research in this direction,
+and offer a robust test-bed for evaluating language models' out-of-distribution
+reasoning capabilities.
+
+## How it works
+
+CLRS-Text is designed as a text-based wrapper around the traces produced by the
+base version of the CLRS benchmark. Therefore, all the key details on how the
+data is generated and how it may be extended to additional distributions and/or
+algorithms are identical to the details for CLRS---please consult the
+[base README file](https://github.com/google-deepmind/clrs/blob/master/README.md)
+for further information.
+
+## What we offer
+
+What we do provide on top of CLRS are the following libraries:
+* `clrs_utils.py`: base functionalities for converting a CLRS trace into
+strings. Note that currently the convertors cannot manage padding, and as a
+result it is recommended to always process traces sampled with a batch size of 1
+and without max-hint tracking.
+* `generate_clrs_text.py`: pre-packaged scripts that can generate the entirety
+of the CLRS-Text training and eval sets with the statistics advertised in our
+paper's evaluation. The outputs are in JSON format. Please see the docstring of
+the script for more details about how to launch it.
+* `huggingface_generators.py`: a convenience function that abstracts away all
+the underlying sampler and generation calls to CLRS and CLRS-Text and provides
+Hugging Face-compatible samples out-of-the-box given a few hyperparameters.
+
+For convenience, we have created a [Hugging Face Collection](https://huggingface.co/collections/tomg-group-umd/clrs-text-668ea13343dab55d6efa1a30)
+which comprises exactly the examples used to train and evaluate our models
+in the original CLRS-Text paper's evaluation.
+
+## Citation
+
+To cite the CLRS-Text Algorithmic Reasoning Language Benchmark:
+
+```latex
+@article{deepmind2024clrstext,
+  title={The CLRS-Text Algorithmic Reasoning Language Benchmark},
+  author={Larisa Markeeva and Sean McLeish and Borja Ibarz and Wilfried Bounsi
+    and Olga Kozlova and Alex Vitvitskyi and Charles Blundell and
+    Tom Goldstein and Avi Schwarzschild and Petar Veli\v{c}kovi\'{c}},
+  journal={arXiv preprint arXiv:2406.04229},
+  year={2024}
+}
+```

clrs/_src/clrs_text/__init__.py

@@ -0,0 +1,18 @@
+# Copyright 2024 DeepMind Technologies Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""The CLRS Text Algorithmic Reasoning Benchmark."""
+
+from clrs._src.clrs_text import clrs_utils
+from clrs._src.clrs_text import huggingface_generators