### Install FlashAttention 3 for Hopper GPUs

Source: https://github.com/dandarm/hrm-ts/blob/main/README.md

Installs FlashAttention version 3, specifically for Hopper GPUs. This involves cloning the repository, navigating to the 'hopper' directory, and running the setup script.

```bash
git clone git@github.com:Dao-AILab/flash-attention.git
cd flash-attention/hopper
python setup.py install
```

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

### Install Python Dependencies

Source: https://github.com/dandarm/hrm-ts/blob/main/README.md

Installs all necessary Python packages listed in the 'requirements.txt' file. This command ensures that all project dependencies are met.

```bash
pip install -r requirements.txt
```

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

### Install CUDA 12.6

Source: https://github.com/dandarm/hrm-ts/blob/main/README.md

Installs CUDA 12.6 using a provided URL and sets the CUDA_HOME environment variable. This is a prerequisite for installing PyTorch with CUDA support and FlashAttention for specific GPU architectures.

```bash
# Install CUDA 12.6
CUDA_URL=https://developer.download.nvidia.com/compute/cuda/12.6.3/local_installers/cuda_12.6.3_560.35.05_linux.run

wget -q --show-progress --progress=bar:force:noscroll -O cuda_installer.run $CUDA_URL
sudo sh cuda_installer.run --silent --toolkit --override

export CUDA_HOME=/usr/local/cuda-12.6
```

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

### Install PyTorch with CUDA 12.6

Source: https://github.com/dandarm/hrm-ts/blob/main/README.md

Installs PyTorch, Torchvision, and Torchaudio with support for CUDA 12.6. This command specifies the index URL for PyTorch wheels compatible with the installed CUDA version.

```bash
# Install PyTorch with CUDA 12.6
PYTORCH_INDEX_URL=https://download.pytorch.org/whl/cu126

pip3 install torch torchvision torchaudio --index-url $PYTORCH_INDEX_URL
```

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

### Hydra Pretraining Configuration (YAML)

Source: https://context7.com/dandarm/hrm-ts/llms.txt

Example Hydra configuration file (cfg_pretrain.yaml) for pretraining. Defines defaults, data paths, training parameters, and learning rate schedules.

```yaml
# config/cfg_pretrain.yaml
defaults:
  - arch: hrm_v1
  - _self_

data_path: data/arc-aug-1000
global_batch_size: 768
epochs: 100000
eval_interval: 10000

lr: 1e-4
lr_min_ratio: 1.0
lr_warmup_steps: 2000

beta1: 0.9
beta2: 0.95
weight_decay: 0.1
puzzle_emb_weight_decay: 0.1
puzzle_emb_lr: 1e-2
```

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

### Install FlashAttention 2 for Ampere or earlier GPUs

Source: https://github.com/dandarm/hrm-ts/blob/main/README.md

Installs FlashAttention version 2 using pip. This is the recommended version for GPUs older than Hopper architecture, such as Ampere.

```bash
pip install flash-attn
```

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

### Define Dataset and Checkpoint Paths

Source: https://github.com/dandarm/hrm-ts/blob/main/arc_eval.ipynb

Specifies the file paths for the dataset and model checkpoints. These constants are crucial for loading training data and pre-trained models. The dataset path can be switched between ARC-1 and ARC-2 configurations.

```python
DATASET_PATH = "data/arc-aug-1000"  # ARC-1
# DATASET_PATH = "data/arc-2-aug-1000"  # ARC-2

CHECKPOINT_PATH = "checkpoints/Arc-aug-1000 ACT-torch/HierarchicalReasoningModel_ACTV1 amphibian-turaco/step_414456"
```

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

### Visualize Input and Ground Truth Data

Source: https://github.com/dandarm/hrm-ts/blob/main/arc_eval.ipynb

This snippet visualizes input and ground truth heatmaps using matplotlib. It iterates through trials, displaying each with a corresponding title and disabling axes. Dependencies include matplotlib and ARC_COLOR_MAP.

```python
if visualize:
                # Show input and ground truth
                axes[fig_id, 0].imshow(global_hmap[input_hash], cmap=ARC_COLOR_MAP)
                axes[fig_id, 0].set_title(f"{name}\nInput")
                axes[fig_id, 0].axis('off')
                
                axes[fig_id, 1].imshow(global_hmap[label_hash], cmap=ARC_COLOR_MAP)
                axes[fig_id, 1].set_title(f"{name}\nAnswer")
                axes[fig_id, 1].axis('off')
                
                trial_id = 2
                for h, stats in p_map[:2]:
                    ans = global_hmap[h]
                    
                    axes[fig_id, trial_id].imshow(ans, cmap=ARC_COLOR_MAP)
                    axes[fig_id, trial_id].set_title(f"{name}\nTrial {trial_id}")
                    axes[fig_id, trial_id].axis('off')
                    
                    trial_id += 1
                
                fig_id += 1
```

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

### Import Libraries for Data Handling and Model Training

Source: https://github.com/dandarm/hrm-ts/blob/main/arc_eval.ipynb

Imports essential Python libraries for file operations, data serialization, pattern matching, hashing, plotting, deep learning (PyTorch), numerical computation (NumPy), and just-in-time compilation (Numba). These are foundational for data processing and model development.

```python
import os
import json
from glob import glob
import hashlib
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors

import torch
import torch.nn.functional as F
import numpy as np
from numba import njit

from dataset.common import inverse_dihedral_transform
```

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

### Lorenz Attractor Time Series Forecasting (PyTorch)

Source: https://context7.com/dandarm/hrm-ts/llms.txt

Complete example for chaotic time series forecasting using the Lorenz system and HRM-TS. Includes data generation, model building, training, and autoregressive inference.

```python
import numpy as np
import torch
from torch.utils.data import DataLoader
from dataset import TimeSeriesWindows
from models.ts_hierarchical_core import TimeSeriesHRMCore
from models.ts_hrm_adapter import TimeSeriesHRM, ts_train_step

def lorenz_series(T=1000, dt=0.01, sigma=10.0, rho=28.0, beta=8/3):
    """Generate Lorenz attractor time series."""
    xyz = np.zeros((T, 3), dtype=np.float32)
    xyz[0] = [1.0, 1.0, 1.0]
    for t in range(1, T):
        x, y, z = xyz[t-1]
        xyz[t] = xyz[t-1] + dt * np.array([
            sigma * (y - x),
            x * (rho - z) - y,
            x * y - beta * z
        ], dtype=np.float32)
    return xyz

# Generate data
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
series = lorenz_series(T=1000)
dataset = TimeSeriesWindows(series, series, T_in=20, T_out=1)
loader = DataLoader(dataset, batch_size=32, shuffle=True)

# Build model
core = TimeSeriesHRMCore(d_model=64, num_heads=4, H_layers=2, L_layers=2)
model = TimeSeriesHRM(core, d_in=3, d_model=64, d_out=3).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# Train
for epoch in range(50):
    for batch in loader:
        batch = {k: v.to(device) for k, v in batch.items()}
        optimizer.zero_grad()
        loss = ts_train_step(model, batch)
        loss.backward()
        optimizer.step()
    if (epoch + 1) % 10 == 0:
        print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

# Autoregressive inference
model.eval()
mu = torch.tensor(dataset.stats["mu"], device=device)
sd = torch.tensor(dataset.stats["sd"], device=device)
x_win = ((torch.tensor(series[:20], device=device) - mu) / sd).unsqueeze(0)

predictions = []
with torch.no_grad():
    for _ in range(980):  # Predict remaining timesteps
        y_hat, _ = model(x_win, T_out=1)
        predictions.append(y_hat[:, -1, :].cpu().numpy())
        y_norm = (y_hat[:, -1:, :] - mu) / sd
        x_win = torch.cat([x_win[:, 1:, :], y_norm], dim=1)

predictions = np.concatenate(predictions, axis=0)
print(f"Predictions shape: {predictions.shape}")  # [980, 3]
```

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

### Execute Test Function

Source: https://github.com/dandarm/hrm-ts/blob/main/arc_eval.ipynb

This snippet demonstrates the execution of a test function, likely for evaluating the model's performance. It calls the 'test' function with the 'visualize' parameter set to False, indicating that no visualization should be generated during this test run.

```python
test(visualize=False)
```

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

### Grid Hashing Utility

Source: https://github.com/dandarm/hrm-ts/blob/main/arc_eval.ipynb

Generates a hash for a given numpy grid. It converts the grid to bytes and includes its shape in the hash to ensure uniqueness for grids with the same content but different dimensions. Input is a numpy array.

```python
def grid_hash(grid: np.ndarray):
    return hash((grid.tobytes(), grid.shape))
```

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

### Load Identifiers and Predictions

Source: https://github.com/dandarm/hrm-ts/blob/main/arc_eval.ipynb

Loads puzzle identifiers and predictions from specified dataset and checkpoint paths. It handles JSON and PyTorch file formats, filters out padding, and concatenates predictions. Dependencies include os, json, glob, torch, and numpy.

```python
def load_identifiers_and_preds(dataset_path: str, checkpoint_path: str):
    # Load puzzle identifiers
    with open(os.path.join(dataset_path, "identifiers.json"), "r") as f:
        identifier_map = json.load(f)
        
    # Load preds
    all_preds = {}
    for filename in glob(f"{checkpoint_path}_all_preds.*"):
        preds = torch.load(filename)
        for k, v in preds.items():
            all_preds.setdefault(k, [])
            all_preds[k].append(v)
            
        del preds

    all_preds = {k: torch.cat(v, dim=0) for k, v in all_preds.items()}
    
    # Remove paddings
    mask = all_preds["puzzle_identifiers"] != PAD_PUZZLE_IDENTIFIER
    all_preds = {k: v[mask] for k, v in all_preds.items()}

    return identifier_map, all_preds
```

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

### Calculate and Print K-Shot Correctness

Source: https://github.com/dandarm/hrm-ts/blob/main/arc_eval.ipynb

This snippet calculates the total correctness for different k-shot values. It iterates through predefined K values, sums up correct predictions, and prints the percentage of correctness for each k. Dependencies include Ks, num_test_correct, tests, and puzzle_labels.

```python
# Total correctness
        for i in range(len(Ks)):
            correct[i] += num_test_correct[i] == len(tests)

    for i, k in enumerate(Ks):
        print (f"{k}-shot: {correct[i] / len(puzzle_labels) * 100:.2f}%")
```

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

### Inverse Augmentation for Grid Transformations

Source: https://github.com/dandarm/hrm-ts/blob/main/arc_eval.ipynb

Applies the inverse of a specified augmentation transformation to a grid. It parses the transformation ID and permutation from the name, calculates the inverse permutation, and applies the inverse dihedral transform. Input is a string representing the augmentation and a numpy array for the grid.

```python
def inverse_aug(name: str, grid: np.ndarray):
    if "_" not in name:
        return grid

    trans_id, perm = name.split("_")[-2:]
    trans_id = int(trans_id[1:])  # Remove "t" letter
    inv_perm = np.argsort(list(perm))
    
    return inv_perm[inverse_dihedral_transform(grid, trans_id)]
```

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

### Define Padding Identifier

Source: https://github.com/dandarm/hrm-ts/blob/main/arc_eval.ipynb

Sets a constant value to represent a padded identifier. This is likely used in data preprocessing or model input to denote padding for sequences or grids.

```python
PAD_PUZZLE_IDENTIFIER = 0
```

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

### Test Function for Model Evaluation

Source: https://github.com/dandarm/hrm-ts/blob/main/arc_eval.ipynb

Evaluates the model's performance by loading data, processing predictions, and comparing them against ground truth labels. It handles augmented and non-augmented puzzles, calculates accuracy based on different K values, and optionally visualizes results. Dependencies include matplotlib.pyplot.

```python
def test(visualize, Ks=[1, 2, 10, 100, 1000]):
    identifier_map, all_preds = load_identifiers_and_preds(DATASET_PATH, CHECKPOINT_PATH)
    
    global_hmap = {}
    
    # Get puzzles and corresponding answers
    puzzle_labels = {}
    for identifier, input, label in zip(all_preds["puzzle_identifiers"], all_preds["inputs"], all_preds["labels"]):
        name = identifier_map[identifier]
        if "_" not in name:   # Not-augmented
            puzzle_labels.setdefault(name, {})
            
            input = crop(input.numpy())
            label = crop(label.numpy())

            input_hash = grid_hash(input)
            label_hash = grid_hash(label)

            global_hmap[input_hash] = input
            global_hmap[label_hash] = label

            assert input_hash not in puzzle_labels[name]
            puzzle_labels[name][input_hash] = label_hash
            
    print ("Number of puzzles", len(puzzle_labels))
    
    # Argmax prediction
    preds = all_preds["logits"].argmax(-1)

    # Collate
    pred_answers = {}
    for identifier, input, pred, q in zip(all_preds["puzzle_identifiers"], all_preds["inputs"], preds, all_preds["q_halt_logits"].sigmoid()):
        name = identifier_map[identifier]
        orig_name = name.split("_")[0]
        
        input = input.numpy()
        input_hash = grid_hash(inverse_aug(name, crop(input)))
        assert input_hash in puzzle_labels[orig_name]
        
        pred = inverse_aug(name, crop(pred.numpy()))
        pred_hash = grid_hash(pred)
        global_hmap[pred_hash] = pred
        
        pred_answers.setdefault(orig_name, {})
        pred_answers[orig_name].setdefault(input_hash, [])
        pred_answers[orig_name][input_hash].append((pred_hash, q.item()))

    # test-1
    if visualize:
        num_figs = sum(len(tests) for name, tests in puzzle_labels.items())
        fig, axes = plt.subplots(num_figs, 4, figsize=(8, num_figs * 4))
        
        fig_id = 0
    
    correct = [0 for _ in range(len(Ks))]
    for name, tests in puzzle_labels.items():
        num_test_correct = [0 for _ in range(len(Ks))]
        for input_hash, label_hash in tests.items():
            p = pred_answers[name][input_hash]
            p_map = {}
            
            for h, q in p:
                p_map.setdefault(h, [0, 0])
                p_map[h][0] += 1
                p_map[h][1] += q
                
            for h, stats in p_map.items():
                stats[1] /= stats[0]
                
            p_map = sorted(p_map.items(), key=lambda kv: kv[1], reverse=True)

            # 2-vote
            for i, k in enumerate(Ks):
                ok = False
                for h, stats in p_map[:k]:
                    ok |= h == label_hash
                    
                num_test_correct[i] += ok

```

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

### ARC Color Map Definition

Source: https://github.com/dandarm/hrm-ts/blob/main/arc_eval.ipynb

Defines a ListedColormap for ARC (Abstraction and Reasoning Corpus) visualization. It assigns specific hexadecimal color codes to symbols 0 through 9, used for representing different elements in the visualizations.

```python
ARC_COLOR_MAP = mcolors.ListedColormap([
    "#000000",  # symbol_0: black
    "#0074D9",  # symbol_1: blue
    "#FF4136",  # symbol_2: red
    "#2ECC40",  # symbol_3: green
    "#FFDC00",  # symbol_4: yellow
    "#AAAAAA",  # symbol_5: grey
    "#F012BE",  # symbol_6: fuschia
    "#FF851B",  # symbol_7: orange
    "#7FDBFF",  # symbol_8: teal
    "#870C25"   # symbol_9: brown
])
```

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

### Crop Grid Function

Source: https://github.com/dandarm/hrm-ts/blob/main/arc_eval.ipynb

Crops a 30x30 grid to the largest possible rectangle that does not contain any EOS tokens (values < 2 or > 11). It iterates through possible row and column sizes to find the maximum area. Requires the numba njit decorator for optimization.

```python
@njit
def crop(grid: np.ndarray):
    # Find maximum-sized rectangle without any EOS token inside.
    grid = grid.reshape(30, 30)

    max_area = 0
    max_size = (0, 0)
    nr, nc = grid.shape
    
    num_c = nc
    for num_r in range(1, nr + 1):
        # Scan for maximum c
        for c in range(1, num_c + 1):
            x = grid[num_r - 1, c - 1]
            if (x < 2) | (x > 11):
                num_c = c - 1
                break
        
        area = num_r * num_c
        if area > max_area:
            max_area = area
            max_size = (num_r, num_c)

    return grid[:max_size[0], :max_size[1]] - 2
```

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

### Slice 1D Typed Array in JavaScript

Source: https://github.com/dandarm/hrm-ts/blob/main/puzzle_visualizer.html

The `slice1D` function extracts a portion of a typed array. It takes the array, a start index, and an end index, returning a new `Uint32Array` containing the specified elements. This is useful for segmenting data, such as input or output sequences for processing.

```javascript
function slice1D(arr, start, end) {
  const result = new Uint32Array(end - start);
  for (let i = start; i < end; i++) {
    result[i - start] = Number(arr[i]);
  }
  return result;
}
```

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

### Launch Training with Hydra Configuration (Bash)

Source: https://context7.com/dandarm/hrm-ts/llms.txt

Commands to launch pretraining using Hydra for configuration management. Shows single GPU, distributed training with torchrun, and overriding architecture parameters.

```bash
# Single GPU training
python pretrain.py data_path=data/arc-aug-1000 epochs=10000 lr=1e-4

# Distributed training with torchrun
torchrun --nproc_per_node=4 pretrain.py \
    data_path=data/arc-aug-1000 \
    global_batch_size=768 \
    epochs=100000 \
    lr=1e-4

# Override architecture parameters
python pretrain.py \
    arch.H_layers=4 \
    arch.L_layers=4 \
    arch.hidden_size=512 \
    arch.num_heads=8
```

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

### Create Sliding Window Datasets with TimeSeriesWindows

Source: https://context7.com/dandarm/hrm-ts/llms.txt

The `TimeSeriesWindows` class preprocesses time series data into sliding input/output windows. It handles automatic normalization and can incorporate optional calendar features. The output is suitable for PyTorch's DataLoader.

```python
import numpy as np
from torch.utils.data import DataLoader
from dataset import TimeSeriesWindows

# Generate sample time series data
np.random.seed(42)
T = 1000
data = np.cumsum(np.random.randn(T, 3), axis=0).astype(np.float32)  # 3-feature series

# Create dataset with 20-step input windows predicting 5 steps ahead
dataset = TimeSeriesWindows(
    X=data,           # Input features [N, D_in]
    Y=data,           # Target values [N, D_out] (can differ from X)
    T_in=20,          # Input window length
    T_out=5,          # Prediction horizon
    stride=1,         # Step between consecutive windows
    stats=None,       # Auto-compute normalization stats from X
    calendar_feats=None  # Optional: callable returning temporal features
)

# Access normalization statistics
print(f"Mean: {dataset.stats['mu'].shape}, Std: {dataset.stats['sd'].shape}")
# Mean: (1, 3), Std: (1, 3)

# Create DataLoader
loader = DataLoader(dataset, batch_size=32, shuffle=True)

# Iterate over batches
for batch in loader:
    x_in = batch["x_in"]      # [32, 20, 3] - normalized input windows
    y_out = batch["y_out"]    # [32, 5, 3] - target values
    print(f"Input shape: {x_in.shape}, Target shape: {y_out.shape}")
    break
```

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

### Run Inference with PyTorch

Source: https://context7.com/dandarm/hrm-ts/llms.txt

This snippet demonstrates how to run inference using a PyTorch model within an inference mode context. It iterates through an evaluation data loader, moves batches to CUDA, initializes model carry, performs a forward pass, and prints the accuracy metric.

```python
with torch.inference_mode():
    for batch in eval_loader:
        batch = {k: v.cuda() for k, v in batch.items()}
        carry = train_state.model.initial_carry(batch)
        carry, _, metrics, preds, _ = train_state.model(carry=carry, batch=batch, return_keys=["logits"])
        print(f"Accuracy: {metrics['accuracy']:.4f}")
```

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

### Evaluate Model Checkpoint (Bash)

Source: https://context7.com/dandarm/hrm-ts/llms.txt

Bash commands to evaluate a trained model checkpoint. Demonstrates loading a specific checkpoint and optionally saving outputs like inputs, labels, and logits.

```bash
# Evaluate a checkpoint
python evaluate.py checkpoint=checkpoints/project/run/step_50000

# With custom output saving
python evaluate.py \
    checkpoint=checkpoints/project/run/step_50000 \
    save_outputs='[inputs,labels,logits]'
```

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

### SwiGLU Feedforward Network (PyTorch)

Source: https://context7.com/dandarm/hrm-ts/llms.txt

Demonstrates the usage of a SwiGLU feedforward network. It takes input tensor 'x' and produces an output tensor of the same shape.

```python
import torch

# Assuming SwiGLU is defined elsewhere
mlp = SwiGLU(hidden_size=64, expansion=4.0)
x = torch.randn(8, 100, 64)
mlp_out = mlp(x)
print(f"MLP output: {mlp_out.shape}")  # [8, 100, 64]
```

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

### Time Series Data Windowing (NumPy/PyTorch)

Source: https://github.com/dandarm/hrm-ts/blob/main/README.md

Demonstrates the data pipeline for time series, which involves windowing input sequences (T_in) to output sequences (T_out) with normalization and optional calendar features. It supports both NumPy and PyTorch backends, selected via dataset initialization.

```python
from dataset import TimeSeriesWindows

# Example usage (conceptual)
dataset_np = TimeSeriesWindows(backend='numpy', ...)
dataset_torch = TimeSeriesWindows(backend='torch', ...)

# The output is PyTorch tensors ready for training.
```

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

### Programmatic Model Evaluation (PyTorch)

Source: https://context7.com/dandarm/hrm-ts/llms.txt

Python code for loading a trained checkpoint and evaluating the model programmatically. Initializes the training state and loads the model weights.

```python
# Programmatic evaluation
import torch
from pretrain import PretrainConfig, init_train_state, create_dataloader

# Load checkpoint
checkpoint_path = "checkpoints/project/run/step_50000"
model_state = torch.load(checkpoint_path, map_location="cuda")

# Initialize model and load weights
train_state = init_train_state(config, train_metadata, world_size=1)
train_state.model.load_state_dict(model_state)
train_state.model.eval()
```

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

### Dynamic Attention Backend Selection and Logging

Source: https://github.com/dandarm/hrm-ts/blob/main/README.md

Implements dynamic selection of attention backends, such as FlashAttention, based on hardware capabilities and configuration. It includes logging functionalities for attention mechanisms via the 'hrm.attn' logger.

```python
from models.layers import set_use_flash
import logging

logger = logging.getLogger('hrm.attn')

# Example usage (conceptual)
# set_use_flash(True) # Enable flash attention if available

# Logs attention mechanism details.
```

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

### RMS Normalization (PyTorch)

Source: https://context7.com/dandarm/hrm-ts/llms.txt

Illustrates the application of RMS normalization to input tensor 'x'. Requires a variance epsilon for numerical stability.

```python
import torch

# Assuming rms_norm is defined elsewhere
x = torch.randn(8, 100, 64)
normalized = rms_norm(x, variance_epsilon=1e-5)
```

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

### Multi-Head Attention Module (PyTorch)

Source: https://context7.com/dandarm/hrm-ts/llms.txt

Shows the instantiation and usage of a multi-head attention module. It takes hidden states and optionally cosine/sine embeddings for attention calculation.

```python
import torch

# Assuming Attention and cos, sin are defined elsewhere
attn = Attention(
    hidden_size=64,
    head_dim=16,
    num_heads=4,
    num_key_value_heads=4,
    causal=False  # Non-causal for bidirectional attention
)
x = torch.randn(8, 100, 64)
attn_out = attn(cos_sin=(cos, sin), hidden_states=x)
print(f"Attention output: {attn_out.shape}")  # [8, 100, 64]
```

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

### Load ARC Dataset from Local Files

Source: https://github.com/dandarm/hrm-ts/blob/main/puzzle_visualizer.html

This asynchronous JavaScript function is triggered when the 'Load' button is clicked. It constructs file paths for inputs, labels, puzzle indices, group indices, and puzzle identifiers based on the selected set and subset. It then calls `loadDataset` to parse these files using Npyjs and FileReader, and finally calls `buildGroupList` to display the loaded data.

```javascript
document.getElementById("loadBtn").addEventListener("click", async () => {
  document.getElementById("groupList").innerHTML = "";
  document.getElementById("puzzleList").innerHTML = "";
  document.getElementById("puzzleView").innerHTML = "";
  const setName = document.getElementById("setSelect").value;
  const subsetName = document.getElementById("subsetSelect").value;
  try {
    await loadDataset(setName, subsetName);
    buildGroupList();
  } catch (err) {
    console.error(err);
    alert("Error while loading dataset: " + err);
  }
});

async function loadDataset(setName, subsetName) {
  const prefix = `${setName}/${subsetName}__`;
  const inputsPath = prefix + "inputs.npy";
  const labelsPath = prefix + "labels.npy";
  const pIdxPath = prefix + "puzzle_indices.npy";
  const gIdxPath = prefix + "group_indices.npy";
  const pIdsPath = prefix + "puzzle_identifiers.npy";
  const identifi
```

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

### Build Group List UI

Source: https://github.com/dandarm/hrm-ts/blob/main/puzzle_visualizer.html

Dynamically generates a list of group elements in the UI based on the global `groupIndicesArr`. Each group is displayed as a clickable span, triggering `onSelectGroup` when clicked.

```javascript
function buildGroupList() {
  document.getElementById("groupList").innerHTML = "<h3>Groups</h3>";
  const groupListDiv = document.getElementById("groupList");
  const nGroups = groupIndicesArr.length - 1;
  for (let g = 0; g < nGroups; g++) {
    const div = document.createElement("span");
    div.className = "group-item";
    div.textContent = `Group ${g}`;
    div.onclick = () => onSelectGroup(g);
    groupListDiv.appendChild(div);
  }
}
```

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

### Apply RoPE to Query/Key Tensors (PyTorch)

Source: https://context7.com/dandarm/hrm-ts/llms.txt

Demonstrates applying Rotary Positional Embeddings (RoPE) to query and key tensors in PyTorch. Supports custom position IDs for variable-length sequences.

```python
import torch

# Assuming apply_rotary_pos_emb, cos, sin are defined elsewhere
q = torch.randn(8, 100, 4, 16)  # [batch, seq, heads, head_dim]
k = torch.randn(8, 100, 4, 16)
q_rot, k_rot = apply_rotary_pos_emb(q, k, cos, sin)

# Custom position IDs for variable-length sequences
position_ids = torch.arange(100)
q_rot, k_rot = apply_rotary_pos_emb(q, k, cos, sin, position_ids=position_ids)
```

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

### Handle Group Selection and Display Puzzles

Source: https://github.com/dandarm/hrm-ts/blob/main/puzzle_visualizer.html

Clears existing puzzle and puzzle view content, then populates the puzzle list with puzzles belonging to the selected group. Each puzzle is displayed as a clickable span, triggering `onSelectPuzzle`.

```javascript
function onSelectGroup(groupIndex) {
  document.getElementById("puzzleList").innerHTML = "";
  document.getElementById("puzzleView").innerHTML = "";
  const puzzleListDiv = document.getElementById("puzzleList");
  puzzleListDiv.innerHTML = `<h4>Puzzles in Group ${groupIndex}</h4>`;
  const firstPuzzle = groupIndicesArr[groupIndex];
  const lastPuzzle = groupIndicesArr[groupIndex + 1];
  for (let p = firstPuzzle; p < lastPuzzle; p++) {
    const puzzleIntId = puzzleIdentifiersArr[p];
    const puzzleStrId = (puzzleIntId < identifiersJson.length) ? identifiersJson[puzzleIntId] : "<unknown>";
    const div = document.createElement("span");
    div.className = "puzzle-item";
    div.textContent = `Puzzle #${p} [ID=${puzzleIntId}: ${puzzleStrId}]`;
    div.onclick = () => onSelectPuzzle(p);
    puzzleListDiv.appendChild(div);
  }
}
```

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

### Read JSON File into Object

Source: https://github.com/dandarm/hrm-ts/blob/main/puzzle_visualizer.html

Reads a local JSON file using the File API, parses its content as JSON, and resolves with the resulting JavaScript object. Includes error handling for file reading and JSON parsing.

```javascript
function readJsonFile(file) {
  return new Promise((resolve, reject) => {
    const reader = new FileReader();
    reader.onload = () => {
      try {
        const obj = JSON.parse(reader.result);
        resolve(obj);
      } catch (err) {
        reject(err);
      }
    };
    reader.onerror = (err) => reject(err);
    reader.readAsText(file);
  });
}
```

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

### Render 2D Grid to Canvas in JavaScript

Source: https://github.com/dandarm/hrm-ts/blob/main/puzzle_visualizer.html

The `renderGrid` function takes a 2D grid and draws it onto an HTML canvas element. It iterates through the grid, using a scaling factor to determine the size of each cell. Each cell's color is determined by the `indexToColor` function based on its value.

```javascript
function renderGrid(grid2d) {
  const rows = grid2d.length;
  const cols = grid2d[0].length;
  const scale = 10;
  const canvas = document.createElement("canvas");
  canvas.width = cols * scale;
  canvas.height = rows * scale;
  canvas.className = "grid-canvas";
  const ctx = canvas.getContext("2d");
  for (let r = 0; r < rows; r++) {
    for (let c = 0; c < cols; c++) {
      const val = grid2d[r][c];
      ctx.fillStyle = indexToColor(val);
      ctx.fillRect(c * scale, r * scale, scale, scale);
    }
  }
  return canvas;
}
```

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

### Handle Local Folder Selection for ARC Dataset

Source: https://github.com/dandarm/hrm-ts/blob/main/puzzle_visualizer.html

This JavaScript function processes the files selected by the user from a local folder. It identifies the 'identifiers.json' file, determines the root directory of the dataset, and reconstructs a `filesByPath` object with relative paths, excluding the root folder. It then enables dataset loading controls.

```javascript
function onFolderSelected(event) {
  filesByPath = {};
  const fileList = event.target.files;
  if (!fileList || fileList.length === 0) {
    alert("No files selected!");
    return;
  }

  const paths = [];
  for (let i = 0; i < fileList.length; i++) {
    const file = fileList[i];
    const relPath = file.webkitRelativePath || file.mozRelativePath || file.name;
    paths.push(relPath);
  }

  const idPath = paths.find(p => p.endsWith("identifiers.json"));
  if (!idPath) {
    alert("Error: No 'identifiers.json' found in the uploaded folder.");
    return;
  }

  let topDir = "";
  const lastSlash = idPath.lastIndexOf("/");
  if (lastSlash >= 0) {
    topDir = idPath.substring(0, lastSlash);
  }

  for (let i = 0; i < fileList.length; i++) {
    const file = fileList[i];
    let relPath = file.webkitRelativePath || file.mozRelativePath || file.name;
    if (topDir && relPath.startsWith(topDir + "/")) {
      relPath = relPath.substring(topDir.length + 1);
    }
    filesByPath[relPath] = file;
  }

  document.getElementById("setSelect").disabled = false;
  document.getElementById("subsetSelect").disabled = false;
  document.getElementById("loadBtn").disabled = false;
}
```

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

### Implement Hierarchical Reasoning Core with TimeSeriesHRMCore

Source: https://context7.com/dandarm/hrm-ts/llms.txt

The `TimeSeriesHRMCore` class implements the core hierarchical reasoning architecture. It consists of alternating high-level and low-level processing modules and utilizes Rotary Position Embeddings (RoPE) for sequence position awareness.

```python
import torch
from models.ts_hierarchical_core import TimeSeriesHRMCore

# Create HRM core with custom configuration
core = TimeSeriesHRMCore(
    d_model=64,                    # Hidden dimension
    num_heads=4,                   # Attention heads
    H_layers=2,                    # High-level module layers
    L_layers=2,                    # Low-level module layers
    H_cycles=1,                    # High-level reasoning cycles
    L_cycles=4,                    # Low-level reasoning cycles per H cycle
    rope_theta=10000.0,            # RoPE base frequency
    max_position_embeddings=2048   # Maximum sequence length
)
```

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

### Map Integer Index to Color in JavaScript

Source: https://github.com/dandarm/hrm-ts/blob/main/puzzle_visualizer.html

The `indexToColor` function maps integer values to specific color strings for rendering. It defines colors for padding (0), end-of-sequence (1), and a palette of 10 colors (2-11) for original data points. Values outside this range default to white.

```javascript
function indexToColor(value) {
  if (value === 0) return "#FFFFFF"; // pad => white
  if (value === 1) return "#DDDDDD"; // eos => light gray
  const colorIdx = value - 2;
  const palette = [
    "#000000", // color0 => black
    "#FF0000", // color1 => red
    "#00FF00", // color2 => green
    "#0000FF", // color3 => blue
    "#FFFF00", // color4 => yellow
    "#FFA500", // color5 => orange
    "#800080", // color6 => purple
    "#00FFFF", // color7 => cyan
    "#FFC0CB", // color8 => pink
    "#808080" // color9 => gray
  ];
  if (colorIdx >= 0 && colorIdx < palette.length) {
    return palette[colorIdx];
  }
  return "#FFFFFF"; // fallback
}
```

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

### TimeSeriesHRM Adapter for Forecasting

Source: https://context7.com/dandarm/hrm-ts/llms.txt

The TimeSeriesHRM class integrates an HRM core with input embedding and a regression head for time series forecasting. It supports both single-step and multi-horizon predictions, as well as deterministic and probabilistic (Gaussian) outputs.

```python
import torch
from models.ts_hierarchical_core import TimeSeriesHRMCore
from models.ts_hrm_adapter import TimeSeriesHRM

# Configuration
d_in = 3      # Input features
d_model = 64  # Hidden dimension
d_out = 3     # Output features

# Build complete model
core = TimeSeriesHRMCore(d_model, num_heads=4, H_layers=2, L_layers=2)
model = TimeSeriesHRM(
    hrm_core=core,
    d_in=d_in,
    d_model=d_model,
    d_out=d_out,
    gaussian=False  # Set True for probabilistic output
)

# Input sequence
x_in = torch.randn(8, 20, 3)  # [batch, T_in, features]

# Single-step forecast
y_pred, log_sigma = model(x_in, T_out=1)
print(f"Prediction shape: {y_pred.shape}")  # [8, 1, 3]

# Multi-horizon forecast (5 steps ahead)
y_pred_multi, _ = model(x_in, T_out=5)
print(f"Multi-horizon shape: {y_pred_multi.shape}")  # [8, 5, 3]

# With probabilistic output
model_prob = TimeSeriesHRM(core, d_in, d_model, d_out, gaussian=True)
y_mean, y_log_sigma = model_prob(x_in, T_out=1)
print(f"Mean: {y_mean.shape}, Log-sigma: {y_log_sigma.shape}")
```

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

### Parse NPY Files with npyjs

Source: https://github.com/dandarm/hrm-ts/blob/main/puzzle_visualizer.html

Reads a File object as an ArrayBuffer, parses it using the npyjs library, and resolves with the parsed data. Handles potential file reading or parsing errors.

```javascript
function parseNpy(file) {
  return new Promise((resolve, reject) => {
    const reader = new FileReader();
    reader.onload = async () => {
      try {
        const arrayBuffer = reader.result;
        const npy = new npyjs();
        resolve(await npy.parse(arrayBuffer));
      } catch (err) {
        reject(err);
      }
    };
    reader.onerror = err => reject(err);
    reader.readAsArrayBuffer(file);
  });
}
```

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

### Time Series HRM Adapter

Source: https://github.com/dandarm/hrm-ts/blob/main/README.md

Orchestrates the integration of time-series specific components with the core Hierarchical Reasoning Model (HRM). The TimeSeriesHRM class acts as an adapter, managing the flow of data through the HRM core and the time-series specific layers.

```python
from models.ts_hrm_adapter import TimeSeriesHRM

# Example usage (conceptual)
time_series_hrm = TimeSeriesHRM(hrm_core=..., ts_embedding=..., ts_head=...)

# This class integrates time-series processing with the HRM core.
```

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

### Time Series Training Step Function

Source: https://context7.com/dandarm/hrm-ts/llms.txt

The `ts_train_step` function facilitates the training process for time series models. It supports various loss functions, including Mean Squared Error (MSE), Huber loss for robustness to outliers, and Gaussian negative log-likelihood for probabilistic models.

```python
import torch
from torch.utils.data import DataLoader
from models.ts_hierarchical_core import TimeSeriesHRMCore
from models.ts_hrm_adapter import TimeSeriesHRM, ts_train_step
from dataset import TimeSeriesWindows
import numpy as np

# Setup data
series = np.random.randn(500, 3).astype(np.float32)
dataset = TimeSeriesWindows(series, series, T_in=20, T_out=1)
loader = DataLoader(dataset, batch_size=32, shuffle=True)

# Setup model
core = TimeSeriesHRMCore(64, num_heads=4)
model = TimeSeriesHRM(core, d_in=3, d_model=64, d_out=3)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# Training loop
for epoch in range(10):
    epoch_loss = 0.0
    for batch in loader:
        optimizer.zero_grad()

        # Standard MSE loss
        loss = ts_train_step(model, batch, gaussian=False, huber=False)

        # Alternative: Huber loss (robust to outliers)
        # loss = ts_train_step(model, batch, huber=True)

        # Alternative: Gaussian NLL (probabilistic model)
        # loss = ts_train_step(model_prob, batch, gaussian=True)

        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()

    print(f"Epoch {epoch+1}, Loss: {epoch_loss/len(loader):.4f}")
```

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

### Time Series Regression Head

Source: https://github.com/dandarm/hrm-ts/blob/main/README.md

Provides the output layer for time-series regression tasks. The TSRegressionHead can operate in either one-shot multi-horizon or autoregressive modes, producing the final predictions.

```python
from models.ts_head import TSRegressionHead

# Example usage (conceptual)
tsr_head = TSRegressionHead(input_dim=..., output_horizon=..., mode='autoregressive')

# Produces predictions in one-shot or autoregressive modes.
```

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

### Time Series Regression Head

Source: https://context7.com/dandarm/hrm-ts/llms.txt

The `TSRegressionHead` maps hidden states from the HRM core to forecast values. It supports both deterministic outputs (predicting a single value) and probabilistic outputs using a Gaussian distribution, providing mean and log-variance.

```python
import torch
from models.ts_head import TSRegressionHead, regression_loss, gaussian_nll

# Deterministic head
head = TSRegressionHead(d_model=64, d_out=3, gaussian=False)
h = torch.randn(8, 5, 64)  # Hidden states [batch, T_out, d_model]

y_pred, log_sigma = head(h)
print(f"Prediction: {y_pred.shape}, Log-sigma: {log_sigma}")
# Prediction: torch.Size([8, 5, 3]), Log-sigma: None

# Probabilistic head
head_prob = TSRegressionHead(d_model=64, d_out=3, gaussian=True)
y_mean, y_log_sigma = head_prob(h)
print(f"Mean: {y_mean.shape}, Log-sigma: {y_log_sigma.shape}")
# Mean: torch.Size([8, 5, 3]), Log-sigma: torch.Size([8, 5, 3])

# Compute losses
target = torch.randn(8, 5, 3)

# MSE loss
mse = regression_loss(y_pred, target, gaussian=False)
print(f"MSE Loss: {mse.item():.4f}")

# Huber loss
huber = regression_loss(y_pred, target, huber=True, delta=1.0)
print(f"Huber Loss: {huber.item():.4f}")

# Gaussian NLL loss
nll = regression_loss(y_mean, target, gaussian=True, log_sigma=y_log_sigma)
print(f"Gaussian NLL: {nll.item():.4f}")

# Direct Gaussian NLL computation
nll_direct = gaussian_nll(y_mean, y_log_sigma, target).mean()
```

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

### Apply Rotary Positional Embeddings (RoPE)

Source: https://github.com/dandarm/hrm-ts/blob/main/README.md

Generalizes Rotary Positional Embeddings (RoPE) to handle sequence windows shorter than the maximum defined length. It correctly indexes cosine/sine values at actual positions and reshapes them for broadcasting.

```python
from models.layers import apply_rotary_pos_emb
import torch

# Example usage (conceptual)
# query = ... # (batch_size, seq_len, num_heads, head_dim)
# key = ...   # (batch_size, seq_len, num_heads, head_dim)
# cos, sin = ... # Precomputed RoPE embeddings

# query, key = apply_rotary_pos_emb(query, key, cos, sin)

# This function handles varying window lengths.
```

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

### Attention with Rotary Position Embeddings (RoPE)

Source: https://context7.com/dandarm/hrm-ts/llms.txt

The `layers` module provides an `Attention` mechanism that incorporates Rotary Position Embeddings (RoPE). It supports efficient computation using either FlashAttention or PyTorch's Scalable Sparse Attention (SDPA) backends. RoPE is used to inject positional information into the attention mechanism.

```python
import torch
from models.layers import (
    Attention, RotaryEmbedding, apply_rotary_pos_emb,
    rms_norm, SwiGLU, CastedLinear
)

# RoPE (Rotary Position Embedding)
rope = RotaryEmbedding(
    dim=16,                        # Head dimension
    max_position_embeddings=512,   # Max sequence length
    base=10000.0                   # Base frequency
)
cos, sin = rope()  # Precomputed cos/sin tables
```

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

### Handle Puzzle Selection and Display Details

Source: https://github.com/dandarm/hrm-ts/blob/main/puzzle_visualizer.html

Clears the puzzle view and displays the selected puzzle's ID and associated identifier. This function is intended to be expanded to show more puzzle details.

```javascript
function onSelectPuzzle(puzzleIndex) {
  const puzzleView = document.getElementById("puzzleView");
  puzzleView.innerHTML = "";
  const puzzleIntId = puzzleIdentifiersArr[puzzleIndex];
  const puzzleStrId = (puzzleIntId < identifiersJson.length) ? identifiersJson[puzzleIntId] : "<unknown>";
  const titleDiv = document.createElement("div");
  titleDiv.className = "puzzle-id";
  titleDiv.textContent = 
```

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

### Embed Time Series with Temporal Encodings using TimeSeriesEmbedding

Source: https://context7.com/dandarm/hrm-ts/llms.txt

The `TimeSeriesEmbedding` module transforms input features and incorporates temporal information. It supports sinusoidal positional encodings or learnable projections from external time features, preparing data for the HRM core.

```python
import torch
from models.ts_embedding import TimeSeriesEmbedding, sinusoidal_pos_encoding

# Create embedding layer with sinusoidal time encodings
embed = TimeSeriesEmbedding(
    d_in=3,           # Number of input features
    d_model=64,       # Model dimension
    d_time=32,        # Size of external time features (if use_sin_time=False)
    use_sin_time=True # Use sinusoidal encodings (default)
)

# Sample input: batch of 8 sequences, 20 timesteps, 3 features
x = torch.randn(8, 20, 3)

# Forward pass
h = embed(x)  # [8, 20, 64] - embedded representation
print(f"Embedded shape: {h.shape}")

# Alternative: use external calendar features
embed_calendar = TimeSeriesEmbedding(d_in=3, d_model=64, d_time=8, use_sin_time=False)
time_features = torch.randn(8, 20, 8)  # e.g., hour, day, month encodings
h_calendar = embed_calendar(x, time_extra=time_features)
print(f"Calendar-embedded shape: {h_calendar.shape}")

# Generate sinusoidal encodings directly
pe = sinusoidal_pos_encoding(T=100, d=64)  # [100, 64]
print(f"Positional encoding shape: {pe.shape}")
```

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

### Temporal Embedding Layer

Source: https://github.com/dandarm/hrm-ts/blob/main/README.md

Implements a temporal embedding layer that projects input features to a specified model dimension (d_model). It incorporates time encodings, including sinusoidal and/or calendar features, using the TimeSeriesEmbedding class.

```python
from models.ts_embedding import TimeSeriesEmbedding

# Example usage (conceptual)
temporal_embedding = TimeSeriesEmbedding(input_dim=..., d_model=..., ...)

# Input features are projected to d_model with time encodings.
```

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

### Decode 1D Sequence to 2D Grid in JavaScript

Source: https://github.com/dandarm/hrm-ts/blob/main/puzzle_visualizer.html

The `decodeGrid` function converts a 1D sequence (flattened array) into a 2D grid. It assumes the sequence length is a perfect square corresponding to `gridSize * gridSize`. The function iterates through the sequence, populating rows and columns of the resulting 2D array.

```javascript
function decodeGrid(seq) {
  const grid = [];
  let idx = 0;
  for (let r = 0; r < gridSize; r++) {
    const row = [];
    for (let c = 0; c < gridSize; c++) {
      row.push(seq[idx]);
      idx++;
    }
    grid.push(row);
  }
  return grid;
}
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.