CLRS Algorithmic Reasoning Benchmark

CLRS represents algorithms as manipulations of sets of objects and their relations using a graph representation. For algorithms involving ordered structures like arrays or rooted trees, ordering is imposed via predecessor links.

This representation allows the benchmark to provide:

• Input/Output pairs: Automatically generated by controlling input distributions to match algorithm preconditions.
• Trajectories of 'hints': These expose the internal state of the algorithm, which can be used to simplify learning or to distinguish between different algorithms that solve the same task (e.g., different sorting methods).

CLRS-Text provides several specialized libraries to facilitate working with textual algorithmic traces:

• clrs_utils.py: Contains base functionalities for converting a CLRS trace into strings. Important: Current converters do not manage padding. It is recommended to process traces sampled with a batch size of 1 and without max-hint tracking.
• generate_clrs_text.py: Pre-packaged scripts to generate the full CLRS-Text training and evaluation sets in JSON format. Refer to the script's docstring for specific launch instructions.
• huggingface_generators.py: A convenience function that abstracts sampler and generation calls to CLRS/CLRS-Text, providing Hugging Face-compatible samples based on provided hyperparameters.

To add a new algorithm to the CLRS suite, follow these three steps:

1. Define the Specification: Determine the input, hint, and output requirements and add them to the SPECS dictionary in clrs/_src/specs.py.
2. Implement the Algorithm: Implement the algorithm in an abstractified form (refer to clrs/_src/algorithms/ for examples).
   • Use probing.push to capture inputs, outputs, and intermediate states (probes).
   • Use probing.finalize to format these probes before returning them with the algorithm output.
3. Implement a Sampler: Create an appropriate input data sampler and add it to the SAMPLERS dictionary in clrs/_src/samplers.py.

Once completed, the algorithm can be accessed via clrs.build_sampler and included in datasets generated by clrs/dataset.py.

You can install the CLRS Algorithmic Reasoning Benchmark via pip from PyPI or directly from the GitHub repository. It is recommended to use a virtual environment to avoid dependency conflicts.

From PyPI:

pip install dm-clrs

From GitHub (latest):

pip install git+https://github.com/google-deepmind/clrs.git

Using a virtual environment:

python3 -m venv clrs_env
source clrs_env/bin/activate
pip install git+https://github.com/google-deepmind/clrs.git

For immediate use, a Hugging Face Collection is available containing the exact examples used to train and evaluate models in the original CLRS-Text paper's evaluation.

Hugging Face Collection

Once installed, you can run the provided example baseline model using the following command. On the first run, the dataset will be automatically downloaded and stored in --dataset_path (defaulting to /tmp/CLRS30).

python3 -m clrs.examples.run

You can load training, evaluation, or test trajectories for a specific algorithm using clrs.create_dataset.

Each item in the dataset is a Feedback namedtuple containing:

• features: A Features namedtuple with inputs, hints (padded to max(T)), and lengths (the true trajectory length for masking).
• outputs: The algorithm's final output.

All fields are ndarrays with a leading batch dimension.

train_ds, num_samples, spec = clrs.create_dataset(
      folder='/tmp/CLRS30',
      algorithm='bfs',
      split='train',
      batch_size=32)

for i, feedback in enumerate(train_ds.as_numpy_iterator()):
  # feedback.features contains inputs, hints, and lengths
  # feedback.outputs contains the target outputs
  if i == 0:
    model.init(feedback.features, initial_seed)
  loss = model.feedback(rng_key, feedback)

If you want to generate samples without using the tensorflow_dataset generator, you can instantiate samplers directly using clrs.build_sampler from clrs/_src/samplers.py.

sampler, spec = clrs.build_sampler(
    name='bfs',
    seed=42,
    num_samples=1000,
    length=16)

def _iterate_sampler(batch_size):
  while True:
    yield sampler.next(batch_size)

for feedback in _iterate_sampler(batch_size=32):
  # Process feedback
  ...