### Set up a virtual environment

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/configuration.md

Recommended setup for isolating project dependencies. Activate the environment before installing the package.

```bash
python3 -m venv clrs_env
source clrs_env/bin/activate  # On Windows: clrs_env\Scripts\activate
pip install dm-clrs
```

--------------------------------

### Install CLRS in a Virtual Environment

Source: https://github.com/google-deepmind/clrs/blob/master/README.md

Set up a virtual environment and install the CLRS benchmark from GitHub within it.

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

--------------------------------

### Run Example Baseline Model

Source: https://github.com/google-deepmind/clrs/blob/master/README.md

Execute the example baseline model provided with the CLRS benchmark. This may trigger dataset download on the first run.

```shell
python3 -m clrs.examples.run
```

--------------------------------

### Install CLRS from GitHub

Source: https://github.com/google-deepmind/clrs/blob/master/README.md

Install the CLRS Algorithmic Reasoning Benchmark directly from GitHub for the latest updates.

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

--------------------------------

### SimpleModel Implementation Example

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/model.md

A concrete implementation of the Model abstract base class. This example demonstrates initializing parameters, generating dummy predictions for outputs, and a placeholder for feedback processing.

```python
from clrs._src import model
from clrs._src import samplers
from clrs._src import specs
from clrs._src import probing
import numpy as np

class SimpleModel(model.Model):
    def __init__(self, spec):
        super().__init__(spec)
        # Initialize model parameters
        self.params = {}
    
    def predict(self, features):
        """Predict outputs and hints from inputs."""
        # features.inputs: tuple of DataPoint
        # features.hints: tuple of DataPoint (reference for training)
        # features.lengths: array of trajectory lengths
        
        predictions = {}
        
        # Implement prediction logic
        for output_name in self._spec[0]:  # Assuming single spec
            stage, loc, type_ = self._spec[0][output_name]
            if stage == specs.Stage.OUTPUT:
                # Generate prediction for this output
                # Shape should match corresponding truth from spec
                dummy_shape = (features.inputs[0].data.shape[0], 10)
                pred_data = np.random.uniform(0, 1, dummy_shape)
                predictions[output_name] = probing.DataPoint(
                    name=output_name,
                    location=loc,
                    type_=type_,
                    data=pred_data
                )
        
        return predictions
    
    def feedback(self, feedback):
        """Update model with ground truth."""
        if feedback is None:
            # Final test phase
            return
        
        # feedback.features: inputs, hints, lengths
        # feedback.outputs: ground truth outputs
        # Update model based on outputs
        pass
```

--------------------------------

### SPECS Dictionary Example for BFS

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/types.md

Example entry from the SPECS dictionary, detailing input, hint, and output data points for the BFS algorithm.

```python
SPECS = {
    'bfs': {
        'pos': ('input', 'node', 'scalar'),
        'key': ('input', 'node', 'scalar'),
        'start_node': ('input', 'node', 'mask_one'),
        'visited': ('hint', 'node', 'mask'),
        'distance': ('output', 'node', 'scalar'),
    },
    'insertion_sort': {
        'pos': ('input', 'node', 'scalar'),
        'key': ('input', 'node', 'scalar'),
        'pred': ('output', 'node', 'should_be_permutation'),
        'pred_h': ('hint', 'node', 'pointer'),
        'i': ('hint', 'node', 'mask_one'),
        'j': ('hint', 'node', 'mask_one'),
    },
    # ... 28 more algorithms ...
}
```

--------------------------------

### predecessor_to_cyclic_predecessor_and_first Example

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/probing.md

Illustrates the conversion of a predecessor array into cyclic predecessor and first-element mask representations using the `predecessor_to_cyclic_predecessor_and_first` function. This example uses a sample predecessor DataPoint.

```python
import clrs
import numpy as np

# Predecessor chain from DFS: 0 -> 1 -> 2 -> None (represented as self-loop)
pred = clrs.DataPoint(
    name='parent',
    location=clrs.Location.NODE,
    type_=clrs.Type.POINTER,
    data=np.array([[2, 2, 2, 0]])  # node 0 has parent 2, etc.
)

cylcic_pred, first = clrs.predecessor_to_cyclic_predecessor_and_first(pred)
print(cyclic_pred)
print(first)
```

--------------------------------

### Install CLRS from PyPI

Source: https://github.com/google-deepmind/clrs/blob/master/README.md

Install the CLRS Algorithmic Reasoning Benchmark using pip from the Python Package Index.

```shell
pip install dm-clrs
```

--------------------------------

### Example: Chaining Sampler Processors

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/samplers.md

Demonstrates how to build a sampler, chain multiple processing functions, and iterate over the processed samples. Requires CLRS and NumPy.

```python
import clrs
import numpy as np

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

def sample_iter():
    while True:
        yield sampler.next(batch_size=32)

# Chain processors
processed = clrs.process_permutations(spec, sample_iter(), enforce_permutations=True)
rng = np.random.RandomState(seed=42)
processed = clrs.process_random_pos(processed, rng)

for feedback in processed:
    # Use processed samples
    pass
```

--------------------------------

### Multi-Algorithm Model Initialization

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/model.md

Example of initializing a model for multiple algorithms (e.g., 'bfs', 'dfs', 'dijkstra') by providing a list of their CLRS specifications.

```python
specs_list = [clrs.SPECS['bfs'], clrs.SPECS['dfs'], clrs.SPECS['dijkstra']]
model = MyModel(spec=specs_list)
```

--------------------------------

### Creating a Model Spec for Multiple Algorithms

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/specs.md

This example shows how to gather specifications for a list of algorithms to be used in a training model. It iterates through the collected specs to print details for each algorithm.

```python
import clrs

# Get specs for algorithms to train on
algorithms = ['bfs', 'dfs', 'dijkstra']
specs_list = [clrs.SPECS[alg] for alg in algorithms]

# Model will need to handle all these specs
for alg_name, spec in zip(algorithms, specs_list):
    print(f"\n{alg_name}:")
    for name, (stage, loc, type_) in spec.items():
        print(f"  {name}: {stage} {loc} {type_}")
```

--------------------------------

### Typical Training Setup with CLRS Sampler

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/configuration.md

Sets up a typical training pipeline including data sampling, model initialization, optimizer configuration, and a basic training loop. Uses `clrs.build_sampler` for data generation.

```python
import clrs
import jax
import jax.numpy as jnp
import optax

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

# Model
model = MyModel(spec=spec)

# Optimizer
optimizer = optax.adam(learning_rate=0.001)

# Training hyperparameters
config = {
    'batch_size': 32,
    'num_steps': 10000,
    'eval_interval': 100,
}

# Training loop
for step in range(config['num_steps']):
    feedback = sampler.next(batch_size=config['batch_size'])
    
    # Forward pass
    predictions = model.predict(feedback.features)
    
    # Evaluate
    scores = clrs.evaluate(feedback.outputs, predictions)
    loss = 1.0 - scores['score']
    
    # Backward pass (implementation-specific)
    
    # Update
    model.feedback(feedback)
    
    if step % config['eval_interval'] == 0:
        print(f"Step {step}: Loss = {loss:.4f}")
```

--------------------------------

### Single-Algorithm Model Initialization

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/model.md

Example of initializing a model for a single algorithm (e.g., 'bfs') using its specific CLRS specification.

```python
spec = clrs.SPECS['bfs']
model = MyModel(spec=spec)
```

--------------------------------

### Iterate Through Dataset and Model Feedback

Source: https://github.com/google-deepmind/clrs/blob/master/README.md

Iterate through the dataset, initialize the model, and provide feedback for training. This example demonstrates the basic training loop.

```python
for i, feedback in enumerate(train_ds.as_numpy_iterator()):
  if i == 0:
    model.init(feedback.features, initial_seed)
  loss = model.feedback(rng_key, feedback)
```

--------------------------------

### Example BFS Algorithm Specification

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/specs.md

An example of a Spec dictionary for the Breadth-First Search (BFS) algorithm. It maps data point names to their stage, location, and type.

```python
bfs_spec = {
    'pos': (Stage.INPUT, Location.NODE, Type.SCALAR),
    'key': (Stage.INPUT, Location.NODE, Type.SCALAR),
    'start_node': (Stage.INPUT, Location.NODE, Type.MASK_ONE),
    'visited': (Stage.HINT, Location.NODE, Type.MASK),
    'distance': (Stage.OUTPUT, Location.NODE, Type.SCALAR),
}
```

--------------------------------

### Example BFS Specification

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/algorithms.md

Provides an example of an algorithm specification, detailing the input, hint, and output data points. This spec defines the structure and types of data associated with an algorithm's execution, such as 'pos', 'key', 'start_node', 'visited', and 'distance' for BFS.

```python
{
    'pos': ('input', 'node', 'scalar'),
    'key': ('input', 'node', 'scalar'),
    'start_node': ('input', 'node', 'mask_one'),
    'visited': ('hint', 'node', 'mask'),
    'distance': ('output', 'node', 'scalar'),
}

```

--------------------------------

### Model Training Loop Example

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/model.md

This snippet demonstrates a typical training loop for a custom model. It shows how to build a sampler, initialize a model, fetch data batches, make predictions, compute loss using clrs.evaluate, and provide feedback to the model for the next iteration.

```python
# Create sampler and model
sampler, spec = clrs.build_sampler('bfs', num_samples=1000, seed=42, length=16)
model = MyModel(spec=spec)

# Training loop
for step in range(num_steps):
    # Get batch
    feedback = sampler.next(batch_size=32)
    
    # Make prediction
    predictions = model.predict(feedback.features)
    
    # Compute loss (using clrs.evaluate)
    scores = clrs.evaluate(feedback.outputs, predictions)
    loss = 1.0 - scores['score']
    
    # Backprop and update (implementation-specific)
    # ... gradient descent ...
    
    # Provide feedback for next step
    model.feedback(feedback)
```

--------------------------------

### DataPoint Example Usage

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/probing.md

Demonstrates the creation of various DataPoint objects for different types of algorithm data, including node-level, edge-level, and graph-level information. Shows how to instantiate DataPoint with NumPy arrays.

```python
import clrs
import numpy as np

# Create data points for algorithm inputs
pos = clrs.DataPoint(
    name='pos',
    location=clrs.Location.NODE,
    type_=clrs.Type.SCALAR,
    data=np.array([[0.1, 0.3, 0.5, 0.7, 0.9]])  # shape: (1, 5)
)

# Node-level data
visited = clrs.DataPoint(
    name='visited',
    location=clrs.Location.NODE,
    type_=clrs.Type.MASK,
    data=np.array([[1, 0, 1, 0, 1]])  # shape: (1, 5)
)

# Edge-level data
adj = clrs.DataPoint(
    name='adjacency',
    location=clrs.Location.EDGE,
    type_=clrs.Type.POINTER,
    data=np.array([[[0, 1, 0, 0, 0],
                     [1, 0, 1, 0, 0],
                     [0, 1, 0, 1, 0],
                     [0, 0, 1, 0, 1],
                     [0, 0, 0, 1, 0]]])  # shape: (1, 5, 5)
)

# Graph-level data
num_nodes = clrs.DataPoint(
    name='num_nodes',
    location=clrs.Location.GRAPH,
    type_=clrs.Type.SCALAR,
    data=np.array([5])
)

print(visited)
# Output: DataPoint(name="visited",	location=node,	type=mask,	data=Array(1, 5))
```

--------------------------------

### Create DataPoint for Node Scalars

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/types.md

Example of creating a DataPoint for node scalars, specifying its name, location, type, and data.

```python
import clrs
import numpy as np

# Input: node scalars for array-based algorithm
node_keys = clrs.DataPoint(
    _name='key',
    _location=clrs.Location.NODE,
    _type_=clrs.Type.SCALAR,
    data=np.array([[5.2, 3.1, 7.8, 1.9]])  # (batch=1, nodes=4)
)
```

--------------------------------

### GAT Processor Example Usage

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/processors.md

Demonstrates how to create and use a GAT processor within a model's forward pass. Ensure input dimensions match the processor's expected configuration.

```python
import clrs.models  # Access processor factory
import haiku as hk
import jax.numpy as jnp

def create_model():
    # Create GAT processor
    gat = clrs.models.processors.GAT(
        out_size=128,
        nb_heads=8,
        activation=jax.nn.relu,
        residual=True
    )
    
    # Use in forward pass
    batch_size, num_nodes, hidden_dim = 32, 16, 128
    node_fts = jnp.ones((batch_size, num_nodes, hidden_dim))
    edge_fts = jnp.ones((batch_size, num_nodes, num_nodes, hidden_dim))
    graph_fts = jnp.ones((batch_size, hidden_dim))
    adj_mat = jnp.ones((batch_size, num_nodes, num_nodes))
    hidden = jnp.ones((batch_size, num_nodes, hidden_dim))
    
    node_out, edge_out = gat(node_fts, edge_fts, graph_fts, adj_mat, hidden)
    return node_out
```

--------------------------------

### Configure JAX Memory Management

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/configuration.md

Sets JAX configurations related to memory management and array modules. This example sets the platform to GPU and specifies JAX NumPy.

```python
import jax

# Preallocate memory
jax.config.update('jax_platform_name', 'gpu')
jax.config.update('jax_array_modules', jax.numpy)

# Disable memory preallocation if needed
jax.config.update('jax_default_prng_impl', 'threefry2x32')
```

--------------------------------

### Processor Integration Pipeline

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/processors.md

Demonstrates how processors are integrated into a baseline model pipeline, typically after encoders and before decoders. This example shows a loop for processing through multiple timesteps.

```python
def create_baseline():
    # Encode inputs
    encoded_nodes = encode_nodes(node_inputs)
    encoded_edges = encode_edges(edge_inputs)
    encoded_graph = encode_graph(graph_inputs)
    
    # Process through multiple timesteps
    hidden = jnp.zeros((batch, nodes, hidden_dim))
    processor = get_processor_factory('gat', out_size=128)()
    
    for t in range(num_steps):
        hidden, _ = processor(
            node_fts=encoded_nodes,
            edge_fts=encoded_edges,
            graph_fts=encoded_graph,
            adj_mat=adjacency,
            hidden=hidden
        )
        
        # Decode at each step
        predictions_t = decode_predictions(hidden)
    
    return predictions_t
```

--------------------------------

### Probing Mechanism Example

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/algorithms.md

Demonstrates how to use the `probing` module to record intermediate states during algorithm execution. This includes recording inputs, intermediate computation steps, and final outputs.

```python
from clrs._src import probing

def example_algorithm(inputs):
    probes = probing.ProbesDict()
    
    # Record input
    probing.push(probes, 'input', 'node', 'scalar', 'key', inputs.key)
    
    # Algorithm execution with recording
    for i in range(len(inputs)):
        # ... computation ...
        probing.push(probes, 'hint', 'node', 'mask_one', 'current_i', mask)
    
    # Record output
    probing.push(probes, 'output', 'node', 'pointer', 'result', output)
    
    # Finalize probes
    finalized = probing.finalize(probes, spec)
    return output, finalized

```

--------------------------------

### CLRS Integration Example: BFS Evaluation

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/evaluation.md

Demonstrates how to build a sampler for Breadth-First Search (BFS), generate feedback, create dummy predictions, and evaluate them using the CLRS library.

```python
import clrs
import numpy as np
import jax
import jax.numpy as jnp

# Sample from dataset
sampler, spec = clrs.build_sampler(
    name='bfs',
    num_samples=100,
    seed=42,
    length=16
)

feedback = sampler.next(batch_size=4)

# Make dummy predictions from model
def make_dummy_predictions(outputs):
    predictions = {}
    for output in outputs:
        # Random predictions with same shape as ground truth
        predictions[output.name] = clrs.DataPoint(
            name=output.name,
            location=output.location,
            type_=output.type_,
            data=np.random.uniform(0, 1, output.data.shape)
        )
    return predictions

predictions = make_dummy_predictions(feedback.outputs)

# Evaluate
scores = clrs.evaluate(feedback.outputs, predictions)
print(f"Overall score: {scores['score']:.3f}")
```

--------------------------------

### Node Data Point Example (Intermediate State)

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/types.md

Defines a node data point representing an intermediate state, like visited nodes over time.

```python
visited = clrs.DataPoint(
    _name='visited',
    _location=clrs.Location.NODE,
    _type_=clrs.Type.MASK,
    data=np.array([[[1, 0, 0, 0],  # t=0
                     [1, 1, 0, 0],  # t=1
                     [1, 1, 1, 0],  # t=2
                     [1, 1, 1, 1]]])  # t=3 (time=4, batch=1, nodes=4)
)
```

--------------------------------

### Dockerfile for CLRS Environment

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/configuration.md

Use this Dockerfile to create a reproducible environment with TensorFlow, JAX, and DM-CLRS installed. It sets the working directory to /workspace.

```dockerfile
FROM tensorflow/tensorflow:latest-gpu

RUN pip install dm-clrs jax jaxlib optax haiku

WORKDIR /workspace
```

--------------------------------

### Build Samplers for Multiple Algorithms

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/algorithms.md

Demonstrates how to create a dictionary of samplers for different algorithms and a model that can handle multiple algorithms. This is useful for training a single model on diverse algorithmic tasks.

```python
samplers_dict = {
    'bfs': clrs.build_sampler('bfs', num_samples=100, seed=42, length=16),
    'dfs': clrs.build_sampler('dfs', num_samples=100, seed=42, length=16),
    'dijkstra': clrs.build_sampler('dijkstra', num_samples=100, seed=42, length=16),
}

specs_list = [clrs.SPECS['bfs'], clrs.SPECS['dfs'], clrs.SPECS['dijkstra']]
model = MyMultiAlgorithmModel(spec=specs_list)
```

--------------------------------

### Build a sampler with BFS algorithm

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/configuration.md

Demonstrates building a sampler for the 'bfs' algorithm, generating 1000 samples with a specific random seed. It also shows optional parameters for tracking steps and truncating decimals.

```python
import clrs

sampler, spec = clrs.build_sampler(
    name='bfs',                          # Algorithm name (required)
    num_samples=1000,                    # Samples to generate (required)
    seed=42,                             # Random seed
    track_max_steps=True,                # Track trajectory length (for on-the-fly)
    truncate_decimals=None,              # Truncate float decimals (algorithm-specific)
    length=16,                           # Max trajectory length (algorithm-specific)
    # Additional algorithm-specific kwargs...
)
```

--------------------------------

### Build Sampler for BFS and Insertion Sort

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/samplers.md

Instantiates samplers for BFS and insertion sort algorithms. Demonstrates pre-generation of samples and on-the-fly sampling.

```python
import clrs

# Pre-generate 1000 samples
sampler, spec = clrs.build_sampler(
    name='bfs',
    num_samples=1000,
    seed=42,
    length=16
)

# Get a batch of 32 samples
feedback = sampler.next(batch_size=32)
inputs = feedback.features.inputs
hints = feedback.features.hints
outputs = feedback.outputs

# On-the-fly sampling (unlimited samples)
sampler_dynamic, spec = clrs.build_sampler(
    name='insertion_sort',
    num_samples=-1,
    seed=42
)

batch = sampler_dynamic.next(batch_size=64)
```

--------------------------------

### Accessing Algorithm Specifications

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/specs.md

Demonstrates how to access specifications for a specific algorithm (e.g., 'bfs') and iterate through its input requirements. It also shows how to list all available algorithms.

```python
import clrs

# Get spec for a specific algorithm
bfs_spec = clrs.SPECS['bfs']

# Check what inputs an algorithm requires
for name, (stage, loc, type_) in bfs_spec.items():
    if stage == clrs.Stage.INPUT:
        print(f"Input: {name}")

# List all algorithms
algorithms = clrs.CLRS_30_ALGS
```

--------------------------------

### get_processor_factory

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/complete-api-index.md

Gets a factory for creating processors. This is used for configuring data processing pipelines.

```APIDOC
## get_processor_factory

### Description
Gets a factory for creating processors. This is used for configuring data processing pipelines.

### Signature
```python
def get_processor_factory(
    processor_name: str,
    out_size: int,
    **kwargs
) -> Callable[[], Processor]
```
```

--------------------------------

### Edge Data Point Example

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/types.md

Defines an edge data point representing adjacency relationships in a graph.

```python
edge_matrix = clrs.DataPoint(
    _name='adjacency',
    _location=clrs.Location.EDGE,
    _type_=clrs.Type.POINTER,
    data=np.array([[[0, 1, 0, 0],
                     [1, 0, 1, 1],
                     [0, 1, 0, 0],
                     [0, 1, 0, 0]]])  # (batch=1, nodes=4, nodes=4)
)
```

--------------------------------

### Get Algorithm Specification

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/README.md

Retrieves the specification for the 'dijkstra' algorithm. This defines the inputs, outputs, and hints for the algorithm.

```python
import clrs
spec = clrs.SPECS['dijkstra']
```

--------------------------------

### Generate Algorithm Traces with BFS Sampler

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/README.md

Demonstrates how to create a sampler for the BFS algorithm and retrieve a batch of feedback data, including inputs, hints, outputs, and lengths.

```python
import clrs

# Create sampler for BFS algorithm
sampler, spec = clrs.build_sampler(
    name='bfs',
    num_samples=1000,
    seed=42,
    length=16
)

# Get a batch of samples
feedback = sampler.next(batch_size=32)

# Access components
inputs = feedback.features.inputs        # Algorithm inputs
hints = feedback.features.hints          # Intermediate states (with time)
outputs = feedback.outputs               # Final outputs
lengths = feedback.features.lengths      # True trajectory lengths
```

--------------------------------

### Get CLRS Folder Name

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/dataset.md

Retrieves the CLRS folder name. This is useful for locating dataset files.

```python
import clrs

folder_name = clrs.get_clrs_folder()
# Returns 'CLRS30_v1.0.0'
```

--------------------------------

### Graph Data Point Example

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/types.md

Defines a graph data point for a single scalar value, such as the number of nodes.

```python
num_nodes = clrs.DataPoint(
    _name='num_nodes',
    _location=clrs.Location.GRAPH,
    _type_=clrs.Type.SCALAR,
    data=np.array([4])  # (batch=1,)
)
```

--------------------------------

### Import Libraries

Source: https://github.com/google-deepmind/clrs/blob/master/clrs/_src/clrs_text/colabs/accuracy_graphs.ipynb

Imports necessary libraries for data manipulation, plotting, and numerical operations. These are foundational for all subsequent code examples.

```python
from matplotlib import pyplot as plt
from matplotlib.ticker import MaxNLocator
import numpy as np
import pandas as pd
import seaborn as sns
```

--------------------------------

### Train Model on Single Algorithm

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/module-reference.md

Demonstrates how to set up a sampler for a single algorithm and train a custom model using the predict, evaluate, and feedback methods.

```python
import clrs
import jax
import jax.numpy as jnp

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

# Create model (subclass of clrs.Model)
model = MyModel(spec=spec)

# Training loop
for step in range(1000):
    feedback = sampler.next(batch_size=32)
    
    # Predict
    predictions = model.predict(feedback.features)
    
    # Evaluate
    scores = clrs.evaluate(feedback.outputs, predictions)
    
    # Update model
    model.feedback(feedback)
```

--------------------------------

### Build Sampler for Dataset Generation

Source: https://github.com/google-deepmind/clrs/blob/master/README.md

Demonstrates how to build a sampler for generating dataset samples without using TensorFlow Datasets. This is useful for custom data generation pipelines.

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

--------------------------------

### Get CLRS Package Version

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/complete-api-index.md

Access the current version of the CLRS package. This is useful for dependency management and ensuring compatibility.

```python
import clrs
print(clrs.__version__)
```

--------------------------------

### Get CLRS Folder Name

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/dataset.md

Retrieves the standard versioned folder name for CLRS30 datasets. This is a simple helper function.

```python
import clrs

folder_name = clrs.get_clrs_folder()
print(folder_name)
```

--------------------------------

### Define FeaturesChunked Namedtuple

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/dataset.md

Defines the FeaturesChunked namedtuple for chunked features, including boundary markers for chunk starts and ends.

```python
import collections

FeaturesChunked = collections.namedtuple(
    'Features', ['inputs', 'hints', 'is_first', 'is_last']
)
```

--------------------------------

### List Available JAX Devices

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/configuration.md

Prints a list of all available computational devices (CPU, GPU, TPU) recognized by JAX.

```python
import jax

# Check available devices
devices = jax.devices()
print(f"Available devices: {devices}")
```

--------------------------------

### Get CLRS Dataset GCP URL

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/dataset.md

Retrieves the Google Cloud Storage URL for the pre-generated CLRS30 dataset. Use this to download the dataset.

```python
import clrs

url = clrs.get_dataset_gcp_url()
# Returns 'https://storage.googleapis.com/dm-clrs/CLRS30_v1.0.0.tar.gz'
```

--------------------------------

### Sampler Usage for BFS

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/samplers.md

Demonstrates how to use the `Sampler` class for the BFS algorithm. Shows instantiation via `build_sampler` and retrieving data batches.

```python
import clrs

# Use build_sampler to instantiate
sampler, spec = clrs.build_sampler('bfs', num_samples=100, seed=42, length=16)

# Get single batch
batch = sampler.next(batch_size=32)

# Get entire pre-generated dataset
all_data = sampler.next(batch_size=None)
```

--------------------------------

### Building Sampler with Unlimited Samples Warning

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/errors.md

Illustrates how to build a sampler with `num_samples=-1`, which triggers a warning about on-the-fly sampling and unlimited samples. This method is memory-efficient but slower.

```python
import clrs

# Use pre-generated samples instead
sampler, spec = clrs.build_sampler(
    name='bfs',
    num_samples=1000,  # Pre-generate 1000 samples
    seed=42,
    length=16
)
```

--------------------------------

### Train a Model with Insertion Sort Sampler

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/README.md

Illustrates the process of creating a sampler for insertion sort, initializing a custom model, and performing a training loop with predictions and feedback.

```python
import clrs

# Create sampler and model
sampler, spec = clrs.build_sampler('insertion_sort', num_samples=1000, seed=42)
model = MyModel(spec=spec)

# Training loop
for step in range(100):
    feedback = sampler.next(batch_size=32)
    
    # Make predictions
    predictions = model.predict(feedback.features)
    
    # Evaluate
    scores = clrs.evaluate(feedback.outputs, predictions)
    
    # Update model
    model.feedback(feedback)
```

--------------------------------

### Using build_sampler for Algorithm Access

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/algorithms.md

Illustrates how to use the `build_sampler` function to obtain a sampler for a specific algorithm. The sampler can then be used to generate batches of algorithm executions and access their traced probes.

```python
import clrs

# Get sampler for algorithm
sampler, spec = clrs.build_sampler(
    name='bfs',
    num_samples=100,
    seed=42,
    length=16
)

# Sampler internally calls the algorithm multiple times with different inputs
feedback = sampler.next(batch_size=32)

# Access traced algorithm execution
inputs = feedback.features.inputs       # Tuple of DataPoints
hints = feedback.features.hints         # Tuple of DataPoints (time series)
lengths = feedback.features.lengths     # True trajectory lengths
outputs = feedback.outputs              # Tuple of DataPoints (final states)

```

--------------------------------

### Get dataset directory from environment variable

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/configuration.md

This Python code retrieves the dataset directory path, defaulting to '/tmp' if the CLRS_DATA_DIR environment variable is not set.

```python
import clrs
import os

dataset_dir = os.getenv('CLRS_DATA_DIR', '/tmp')
```

--------------------------------

### On-the-fly Sample Generation with BFS Sampler

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/configuration.md

Configure the sampler to generate samples dynamically as needed, which is more memory-efficient. `num_samples=-1` indicates on-the-fly generation.

```python
sampler_memory = clrs.build_sampler(
    name='bfs',
    num_samples=-1,        # Generate on-the-fly
    seed=42,
    track_max_steps=False  # Don't track max steps
)
```

--------------------------------

### Iterate Sampler for Batched Data

Source: https://github.com/google-deepmind/clrs/blob/master/README.md

Provides an example of an iterator function to continuously yield batches of data from a sampler. This is typically used within a training loop.

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

for feedback in _iterate_sampler(batch_size=32):
  ...
```

--------------------------------

### Get Processor Factory

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/processors.md

Use the `get_processor_factory` function to obtain a callable that constructs processor instances. Specify the processor name, output size, and any processor-specific hyperparameters.

```python
import clrs

# Get a processor factory
processor_factory = clrs.get_processor_factory(
    processor_name='gat',
    out_size=128,
    nb_heads=8
)

# Create processor in Haiku transformed function
def forward_pass():
    processor = processor_factory()
    # Use processor...
```

--------------------------------

### FeaturesChunked Type Definition

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/complete-api-index.md

Defines the `FeaturesChunked` named tuple, used for chunked data. It includes inputs, hints, and flags indicating the start and end of chunks.

```python
FeaturesChunked = namedtuple('Features', ['inputs', 'hints', 'is_first', 'is_last'])
# inputs: Tuple[DataPoint, ...]
# hints: Tuple[DataPoint, ...]
# is_first: np.ndarray
# is_last: np.ndarray
```

--------------------------------

### Create dataset with automatic download

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/configuration.md

This Python code snippet demonstrates how to create a dataset, which will automatically download if not found in the specified folder. It configures the dataset for Breadth-First Search (BFS) with a batch size of 32.

```python
import clrs

train_ds, num_samples, spec = clrs.create_dataset(
    folder='/tmp/CLRS30',  # Download location
    algorithm='bfs',
    split='train',
    batch_size=32
)
# First run: downloads ~1GB from GCP
```

--------------------------------

### Building Sampler and Ignoring Kwargs Warning

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/errors.md

Demonstrates building a sampler and highlights the warning that occurs when passing algorithm-specific keyword arguments that the sampler does not accept. It suggests checking the sampler's signature.

```python
import clrs
import inspect

# Build sampler to see what kwargs it accepts
sampler, spec = clrs.build_sampler(
    name='bfs',
    num_samples=100,
    seed=42,
    length=16
    # Only pass valid kwargs for bfs sampler
)
```

--------------------------------

### Training Loop with Algorithm Index Input

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/algorithms.md

Illustrates a training loop where samplers are used to generate feedback, algorithm indices are added to the input features for routing, and predictions are made and evaluated. This shows a practical application of the multi-algorithm model.

```python
# Training loop
for algo_name, sampler in samplers_dict.items():
    feedback = sampler.next(batch_size=32)
    
    # Add algorithm index to inputs for routing
    algo_idx = clrs.DataPoint(
        name=clrs.ALGO_IDX_INPUT_NAME,
        location=clrs.Location.GRAPH,
        type_=clrs.Type.CATEGORICAL,
        data=np.array([list(samplers_dict.keys()).index(algo_name)])
    )
    
    predictions = model.predict(feedback.features)
    scores = clrs.evaluate(feedback.outputs, predictions)
```

--------------------------------

### Define Spec Type Alias with Example

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/types.md

Defines Spec as a type alias for an algorithm specification dictionary. Maps data point names to their stage, location, and type classifications.

```python
Spec = Dict[str, Tuple[str, str, str]]
```

```python
{
    'point_name': (stage, location, type),
    'pos': ('input', 'node', 'scalar'),
    'visited': ('hint', 'node', 'mask'),
    'distance': ('output', 'node', 'scalar'),
    ...
}
```

--------------------------------

### Pre-generate Samples

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/README.md

Builds a sampler for a specified algorithm and number of samples. Pre-generating samples can speed up the training process.

```python
sampler, spec = clrs.build_sampler('bfs', num_samples=1000)
```

--------------------------------

### Use Pre-Generated Datasets

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/module-reference.md

Shows how to load and iterate through pre-generated datasets for training, accessing features and outputs for model training.

```python
import clrs

# Download and extract CLRS30_v1.0.0.tar.gz to /tmp/CLRS30

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

for feedback in train_ds.take(100):
    # feedback is Feedback namedtuple
    inputs = feedback.features.inputs
    hints = feedback.features.hints
    lengths = feedback.features.lengths
    outputs = feedback.outputs
    
    # Train model...
```

--------------------------------

### Pre-generate Samples with BFS Sampler

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/configuration.md

Use this configuration to pre-generate a fixed number of samples for faster access. Setting `track_max_steps=True` can help avoid recompilation.

```python
sampler_fast = clrs.build_sampler(
    name='bfs',
    num_samples=1000,      # Pre-generate 1000 samples
    seed=42,
    track_max_steps=True   # Avoid recompilation
)
```

--------------------------------

### predecessor_to_cyclic_predecessor_and_first

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/probing.md

Converts a predecessor array DataPoint into a cyclic predecessor representation and a binary mask indicating chain starts. This transformation is useful for processing chain-like structures in neural networks.

```APIDOC
## predecessor_to_cyclic_predecessor_and_first

### Description
Converts a predecessor array DataPoint into a cyclic predecessor representation and a binary mask indicating chain starts. This transformation is useful for processing chain-like structures in neural networks.

### Parameters
- **pred** (DataPoint): Pointer-type data point with predecessor indices.

### Returns
- **Tuple[DataPoint, DataPoint]**: A tuple containing the cyclic predecessor DataPoint and the first node mask DataPoint.
```

--------------------------------

### OutputClass Enum for Masked Data

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/specs.md

Defines special values for masked outputs: POSITIVE (1), NEGATIVE (0), and MASKED (-1). Use these to indicate valid examples or masked-out regions.

```python
class OutputClass:
    POSITIVE = 1
    NEGATIVE = 0
    MASKED = -1
```

--------------------------------

### Configure CLRS30 benchmark samplers

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/configuration.md

Sets up training and testing samplers using the fixed parameters defined in `clrs.CLRS30_CONFIG`. This ensures consistent benchmark evaluation.

```python
import clrs

CLRS30_CONFIG = clrs.CLRS30
# {
#     'train': {
#         'num_samples': 1000,
#         'length': 16,
#         'seed': 1,
#     },
#     'val': {
#         'num_samples': 32,
#         'length': 16,
#         'seed': 2,
#     },
#     'test': {
#         'num_samples': 32,
#         'length': 64,
#         'seed': 3,
#     },
# }

# Use these settings:
sampler_train, spec = clrs.build_sampler(
    name='bfs',
    num_samples=CLRS30_CONFIG['train']['num_samples'],
    seed=CLRS30_CONFIG['train']['seed'],
    length=CLRS30_CONFIG['train']['length']
)

sampler_test, _ = clrs.build_sampler(
    name='bfs',
    num_samples=CLRS30_CONFIG['test']['num_samples'],
    seed=CLRS30_CONFIG['test']['seed'],
    length=CLRS30_CONFIG['test']['length']
)
```

--------------------------------

### Train on Multiple Algorithms

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/module-reference.md

Illustrates training a model on multiple algorithms by creating separate samplers for each and adapting the model to handle multiple specifications.

```python
import clrs

algorithms = ['bfs', 'dfs', 'dijkstra']
samplers_dict = {
    alg: clrs.build_sampler(alg, num_samples=1000, seed=42, length=16)
    for alg in algorithms
}

specs_list = [clrs.SPECS[alg] for alg in algorithms]
model = MyMultiAlgorithmModel(spec=specs_list)

# Training loop
for algo in algorithms:
    sampler = samplers_dict[algo]
    feedback = sampler.next(batch_size=32)
    
    # Add algorithm index
    algo_idx = clrs.DataPoint(
        name=clrs.ALGO_IDX_INPUT_NAME,
        location=clrs.Location.GRAPH,
        type_=clrs.Type.CATEGORICAL,
        data=np.array([algorithms.index(algo)])
    )
    
    predictions = model.predict(feedback.features)
    scores = clrs.evaluate(feedback.outputs, predictions)
    model.feedback(feedback)
```

--------------------------------

### Multi-Algorithm Model with Dynamic Algorithm Index

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/model.md

Demonstrates how to include an algorithm index as an input feature for multi-algorithm models, allowing the model to differentiate between algorithms during processing.

```python
# Or dynamically add algorithm index as input
algo_idx = clrs.DataPoint(
    name=clrs.ALGO_IDX_INPUT_NAME,
    location=clrs.Location.GRAPH,
    type_=clrs.Type.CATEGORICAL,
    data=np.array([0])  # 0 for bfs, 1 for dfs, etc.
)
```

--------------------------------

### Load Pre-Generated Dataset

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/README.md

Loads a pre-generated dataset for a specified algorithm ('bfs') and split ('train') with a given batch size. Returns the training dataset, number of samples, and algorithm specification.

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

--------------------------------

### Create Processor Factory

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/configuration.md

Builds a processor factory for a specified algorithm. Accepts algorithm names like 'gat', 'mpnn', 'pgn', etc.

```python
import clrs

processor_factory = clrs.get_processor_factory(
    processor_name='gat',   # 'gat', 'gat_v2', 'mpnn', 'pgn', 'deepsets', 'memnet'
    out_size=128,
    nb_heads=8
)
```

--------------------------------

### Create dataset from a custom location

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/configuration.md

Configures dataset creation using a custom directory path. It retrieves the correct folder name using `clrs.get_clrs_folder()` and then creates the dataset for BFS.

```python
import clrs

dataset_dir = '/path/to/datasets'

# Get correct folder name
folder_name = clrs.get_clrs_folder()  # Returns 'CLRS30_v1.0.0'

train_ds, num_samples, spec = clrs.create_dataset(
    folder=f'{dataset_dir}/{folder_name}',
    algorithm='bfs',
    split='train',
    batch_size=32
)
```

--------------------------------

### Accessing Algorithms Implementation

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/complete-api-index.md

Demonstrates how to import and access specific algorithm implementations from the CLRS library. These are typically used internally by samplers.

```python
import clrs
from clrs._src import algorithms

# These are used internally by samplers, not directly
bfs_impl = algorithms.bfs
```

--------------------------------

### Typical Training Loop Structure

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/model.md

Illustrates the basic structure of a training loop for CLRS models, showing the import of the CLRS library.

```python
import clrs
```

--------------------------------

### Typical Decoder Usage in a Complete Model

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/api-reference/decoders.md

Demonstrates how to construct and apply decoder layers within a model's forward pass, including post-processing for specific output types like PERMUTATION_POINTER. Requires Haiku, JAX, and CLRS libraries.

```python
import haiku as hk
import jax.numpy as jnp
from clrs._src import decoders, specs, samplers

def model_forward(features: samplers.Features, spec: specs.Spec):
    # ... encoder and processor steps ...
    hidden = jnp.ones((32, 16, 128))  # (batch, nodes, hidden_dim)
    
    predictions = {}
    
    for name, (stage, loc, type_) in spec.items():
        if stage == specs.Stage.OUTPUT:
            # Get decoder for this output
            decoder_layers = decoders.construct_decoders(
                loc=loc,
                t=type_,
                hidden_dim=hidden.shape[-1],
                nb_dims=10,
                name=f'decoder_{name}'
            )
            
            # Apply decoder layers
            output = hidden
            for layer in decoder_layers:
                output = layer(output)
            
            # Post-process if necessary
            if type_ == specs.Type.PERMUTATION_POINTER:
                # Convert to doubly stochastic
                output = decoders.log_sinkhorn(
                    output.squeeze(-1),
                    steps=10,
                    temperature=0.1,
                    zero_diagonal=False,
                    noise_rng_key=None
                )
            
            predictions[name] = output
    
    return predictions
```

--------------------------------

### CLRS_30_ALGS List

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/types.md

List of all 30 algorithms with canonical names for use in building samplers and specifications.

```python
CLRS_30_ALGS = [
    'articulation_points', 'activity_selector', 'bellman_ford', 'bfs',
    'binary_search', 'bridges', 'bubble_sort', 'dag_shortest_paths', 'dfs',
    'dijkstra', 'find_maximum_subarray_kadane', 'floyd_warshall', 'graham_scan',
    'heapsort', 'insertion_sort', 'jarvis_march', 'kmp_matcher', 'lcs_length',
    'matrix_chain_order', 'minimum', 'mst_kruskal', 'mst_prim',
    'naive_string_matcher', 'optimal_bst', 'quickselect', 'quicksort',
    'segments_intersect', 'strongly_connected_components', 'task_scheduling',
    'topological_sort'
]
```

--------------------------------

### Define Multi-Algorithm Model

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/configuration.md

Initializes a custom model that can handle multiple algorithms. Pass a list of specification objects, one for each algorithm.

```python
import clrs

algorithms = ['bfs', 'dfs', 'dijkstra']
specs_list = [clrs.SPECS[alg] for alg in algorithms]

my_model = MyModel(spec=specs_list)
```

--------------------------------

### Build and Run CLRS Docker Container

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/configuration.md

Build the Docker image with the tag 'clrs:latest' and then run an interactive container with GPU access enabled.

```bash
docker build -t clrs:latest .
docker run --gpus all -it clrs:latest
```

--------------------------------

### create_dataset

Source: https://github.com/google-deepmind/clrs/blob/master/_autodocs/complete-api-index.md

Creates a dataset from a specified folder for a given algorithm and split. This function is useful for preparing data for training or evaluation.

```APIDOC
## create_dataset

### Description
Creates a dataset from a specified folder for a given algorithm and split. This function is useful for preparing data for training or evaluation.

### Signature
```python
def create_dataset(
    folder: str,
    algorithm: str,
    split: str,
    batch_size: int
) -> Tuple[tf.data.Dataset, int, specs.Spec]
```
```