Package: dm-clrs
The CLRS Algorithmic Reasoning BenchmarkLearning representations of algorithms is an emerging area of machine learning,
seeking to bridge concepts from neural networks with classical algorithms. The
CLRS Algorithmic Reasoning Benchmark (CLRS) consolidates and extends previous
work torward evaluation algorithmic reasoning by providing a suite of
implementations of classical algorithms. These algorithms have been selected
from the third edition of the standardIntroduction to Algorithmsby Cormen,
Leiserson, Rivest and Stein.Getting startedThe CLRS Algorithmic Reasoning Benchmark can be installed with pip, either from
PyPI:pipinstalldm-clrsor directly from GitHub (updated more frequently):pipinstallgit+git://github.com/deepmind/clrs.gitYou may prefer to install it in a virtual environment if any requirements
clash with your Python installation:python3-mvenvclrs_envsourceclrs_env/bin/activate
pipinstallgit+git://github.com/deepmind/clrs.gitOnce installed you can run our example baseline model:python3-mclrs.examples.runIf this is the first run of the example, the dataset will be downloaded and
stored in--dataset_path(default '/tmp/CLRS30').
Alternatively, you can also download and extracthttps://storage.googleapis.com/dm-clrs/CLRS30_v1.0.0.tar.gzAlgorithms as graphsCLRS implements the selected algorithms in an idiomatic way, which aligns as
closely as possible to the original CLRS 3ed pseudocode. By controlling the
input data distribution to conform to the preconditions we are able to
automatically generate input/output pairs. We additionally provide trajectories
of "hints" that expose the internal state of each algorithm, to both optionally
simplify the learning challenge and to distinguish between different algorithms
that solve the same overall task (e.g. sorting).In the most generic sense, algorithms can be seen as manipulating sets of
objects, along with any relations between them (which can themselves be
decomposed into binary relations). Accordingly, we study all of the algorithms
in this benchmark using a graph representation. In the event that objects obey a
more strict ordered structure (e.g. arrays or rooted trees), we impose this
ordering through inclusion of predecessor links.How it worksFor each algorithm, we provide a canonical set oftrain,evalandtesttrajectories for benchmarking out-of-distribution generalization.TrajectoriesProblem SizeTrain100016Eval32 x multiplier16Test32 x multiplier64Here, "problem size" refers to e.g. the length of an array or number of nodes in
a graph, depending on the algorithm. "multiplier" is an algorithm-specific
factor that increases the number of availableevalandtesttrajectories
to compensate for paucity of evaluation signals. "multiplier" is 1 for all
algorithms except:Maximum subarray (Kadane), for which "multiplier" is 32.Quick select, minimum, binary search, string matchers (both naive and KMP),
and segment intersection, for which "multiplier" is 64.The trajectories can be used like so:train_ds,num_samples,spec=clrs.create_dataset(folder='/tmp/CLRS30',algorithm='bfs',split='train',batch_size=32)fori,feedbackinenumerate(train_ds.as_numpy_iterator()):ifi==0:model.init(feedback.features,initial_seed)loss=model.feedback(rng_key,feedback.features)Here,feedbackis anamedtuplewith the following structure:Feedback=collections.namedtuple('Feedback',['features','outputs'])Features=collections.namedtuple('Features',['inputs','hints','lengths'])where the content ofFeaturescan be used for training andoutputsis
reserved for evaluation. Each field of the tuple is anndarraywith a leading
batch dimension. Becausehintsare provided for the full algorithm trajectory,
these contain an additional time dimension padded up to the maximum lengthmax(T)of any trajectory within the dataset. Thelengthsfield specifies the
true lengtht <= max(T)for each trajectory, which can be used e.g. for loss
masking.Please see theexamplesdirectory for full working Graph Neural Network (GNN)
examples using JAX and the DeepMind JAX Ecosystem of libraries.What we provideAlgorithmsOur initial CLRS-30 benchmark includes the following 30 algorithms. We aim to
support more algorithms in the future.SortingInsertion sortBubble sortHeapsort (Williams, 1964)Quicksort (Hoare, 1962)SearchingMinimumBinary searchQuickselect (Hoare, 1961)Divide and conquerMaximum subarray (Kadane's variant) (Bentley, 1984)GreedyActivity selection (Gavril, 1972)Task scheduling (Lawler, 1985)Dynamic programmingMatrix chain multiplicationLongest common subsequenceOptimal binary search tree (Aho et al., 1974)GraphsDepth-first search (Moore, 1959)Breadth-first search (Moore, 1959)Topological sorting (Knuth, 1973)Articulation pointsBridgesKosaraju's strongly connected components algorithm (Aho et al., 1974)Kruskal's minimum spanning tree algorithm (Kruskal, 1956)Prim's minimum spanning tree algorithm (Prim, 1957)Bellman-Ford algorithm for single-source shortest paths (Bellman, 1958)Dijkstra's algorithm for single-source shortest paths (Dijkstra et al., 1959)Directed acyclic graph single-source shortest pathsFloyd-Warshall algorithm for all-pairs shortest-paths (Floyd, 1962)StringsNaïve string matchingKnuth-Morris-Pratt (KMP) string matcher (Knuth et al., 1977)GeometrySegment intersectionGraham scan convex hull algorithm (Graham, 1972)Jarvis' march convex hull algorithm (Jarvis, 1973)BaselinesModels consist of aprocessorand a number ofencodersanddecoders.
We provide JAX implementations of the following GNN baseline processors:Deep Sets (Zaheer et al., NIPS 2017)End-to-End Memory Networks (Sukhbaatar et al., NIPS 2015)Graph Attention Networks (Veličković et al., ICLR 2018)Graph Attention Networks v2 (Brody et al., ICLR 2022)Message-Passing Neural Networks (Gilmer et al., ICML 2017)Pointer Graph Networks (Veličković et al., NeurIPS 2020)If you want to implement a new processor, the easiest way is to add
it in theprocessors.pyfile and make it available through theget_processor_factorymethod there. A processor should have a__call__method like this:__call__(self,
         node_fts, edge_fts, graph_fts,
         adj_mat, hidden,
         nb_nodes, batch_size)wherenode_fts,edge_ftsandgraph_ftswill be float arrays of shapebatch_sizexnb_nodesx H,batch_sizexnb_nodesxnb_nodesx H,
andbatch_sizex H with encoded features for
nodes, edges and graph respectively,adj_matabatch_sizexnb_nodesxnb_nodesboolean
array of connectivity built from hints and inputs, andhiddenabatch_sizexnb_nodesx H float array with the previous-step outputs
of the processor. The method should return abatch_sizexnb_nodesx H
float array.For more fundamentally different baselines, it is necessary to create a new
class that extends the Model API (as found withinclrs/_src/model.py).clrs/_src/baselines.pyprovides one example of how this can be done.Creating your own datasetWe provide atensorflow_datasetgenerator class indataset.py. This file can
be modified to generate different versions of the available algorithms, and it
can be built by usingtfds buildafter following the installation instructions
athttps://www.tensorflow.org/datasets.Alternatively, you can generate samples without going throughtfdsby
instantiating samplers with thebuild_samplermethod inclrs/_src/samplers.py, like so: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):
  ...Adding new algorithmsAdding a new algorithm to the task suite requires the following steps:Determine the input/hint/output specification of your algorithm, and include
it within theSPECSdictionary ofclrs/_src/specs.py.Implement the desired algorithm in an abstractified form. Examples of this
can be found throughout theclrs/_src/algorithms/folder.Choose appropriate moments within the algorithm’s execution to create probes
that capture the inputs, outputs and all intermediate state (using
theprobing.pushfunction).Once generated, probes must be formatted using theprobing.finalizemethod, and should be returned together with the algorithm output.Implement an appropriate input data sampler for your algorithm,
and include it in theSAMPLERSdictionary withinclrs/_src/samplers.py.Once the algorithm has been added in this way, it can be accessed with thebuild_samplermethod, and will also be incorporated to the dataset if
regenerated with the generator class indataset.py, as described above.CitationTo cite the CLRS Algorithmic Reasoning Benchmark:@article{deepmind2022clrs,
  title={The CLRS Algorithmic Reasoning Benchmark},
  author={Petar Veli\v{c}kovi\'{c}and Adri\`{a}Puigdom\`{e}nech Badia and
    David Budden and Razvan Pascanu and Andrea Banino and Misha Dashevskiy and
    Raia Hadsell and Charles Blundell},
  journal={arXiv preprint arXiv:2205.15659},
  year={2022}}