SYMBOL INDEX (109 symbols across 14 files)

FILE: assets/npyjs.js
  class npyjs (line 1) | class npyjs {
    method constructor (line 3) | constructor(opts) {
    method float16ToFloat32Array (line 78) | float16ToFloat32Array(float16Array) {
    method float16ToFloat32 (line 89) | static float16ToFloat32(float16) {
    method parse (line 116) | parse(arrayBufferContents) {
    method load (line 155) | async load(filename, callback, fetchArgs) {

FILE: dataset/build_arc_dataset.py
  class DataProcessConfig (line 19) | class DataProcessConfig(BaseModel):
  class ARCPuzzle (line 37) | class ARCPuzzle:
  function arc_grid_to_np (line 43) | def arc_grid_to_np(grid: List[List[int]]):
  function np_grid_to_seq_translational_augment (line 54) | def np_grid_to_seq_translational_augment(inp: np.ndarray, out: np.ndarra...
  function puzzle_hash (line 81) | def puzzle_hash(puzzle: dict):
  function convert_single_arc_puzzle (line 98) | def convert_single_arc_puzzle(results: dict, default_name: str, puzzle: ...
  function load_puzzles_arcagi (line 148) | def load_puzzles_arcagi(results: dict, dataset_path: str, config: DataPr...
  function convert_dataset (line 184) | def convert_dataset(config: DataProcessConfig):
  function main (line 286) | def main(config: DataProcessConfig):

FILE: dataset/build_maze_dataset.py
  class DataProcessConfig (line 22) | class DataProcessConfig(BaseModel):
  function convert_subset (line 30) | def convert_subset(set_name: str, config: DataProcessConfig):
  function preprocess_data (line 136) | def preprocess_data(config: DataProcessConfig):

FILE: dataset/build_sudoku_dataset.py
  class DataProcessConfig (line 18) | class DataProcessConfig(BaseModel):
  function shuffle_sudoku (line 27) | def shuffle_sudoku(board: np.ndarray, solution: np.ndarray):
  function convert_subset (line 60) | def convert_subset(set_name: str, config: DataProcessConfig):
  function preprocess_data (line 163) | def preprocess_data(config: DataProcessConfig):

FILE: dataset/common.py
  class PuzzleDatasetMetadata (line 12) | class PuzzleDatasetMetadata(pydantic.BaseModel):
  function dihedral_transform (line 27) | def dihedral_transform(arr: np.ndarray, tid: int) -> np.ndarray:
  function inverse_dihedral_transform (line 50) | def inverse_dihedral_transform(arr: np.ndarray, tid: int) -> np.ndarray:

FILE: evaluate.py
  class EvalConfig (line 13) | class EvalConfig(pydantic.BaseModel):
  function launch (line 19) | def launch():

FILE: models/common.py
  function trunc_normal_init_ (line 7) | def trunc_normal_init_(tensor: torch.Tensor, std: float = 1.0, lower: fl...

FILE: models/hrm/hrm_act_v1.py
  class HierarchicalReasoningModel_ACTV1InnerCarry (line 16) | class HierarchicalReasoningModel_ACTV1InnerCarry:
  class HierarchicalReasoningModel_ACTV1Carry (line 22) | class HierarchicalReasoningModel_ACTV1Carry:
  class HierarchicalReasoningModel_ACTV1Config (line 31) | class HierarchicalReasoningModel_ACTV1Config(BaseModel):
  class HierarchicalReasoningModel_ACTV1Block (line 60) | class HierarchicalReasoningModel_ACTV1Block(nn.Module):
    method __init__ (line 61) | def __init__(self, config: HierarchicalReasoningModel_ACTV1Config) -> ...
    method forward (line 77) | def forward(self, cos_sin: CosSin, hidden_states: torch.Tensor) -> tor...
  class HierarchicalReasoningModel_ACTV1ReasoningModule (line 86) | class HierarchicalReasoningModel_ACTV1ReasoningModule(nn.Module):
    method __init__ (line 87) | def __init__(self, layers: List[HierarchicalReasoningModel_ACTV1Block]):
    method forward (line 92) | def forward(self, hidden_states: torch.Tensor, input_injection: torch....
  class HierarchicalReasoningModel_ACTV1_Inner (line 102) | class HierarchicalReasoningModel_ACTV1_Inner(nn.Module):
    method __init__ (line 103) | def __init__(self, config: HierarchicalReasoningModel_ACTV1Config) -> ...
    method _input_embeddings (line 146) | def _input_embeddings(self, input: torch.Tensor, puzzle_identifiers: t...
    method empty_carry (line 168) | def empty_carry(self, batch_size: int):
    method reset_carry (line 174) | def reset_carry(self, reset_flag: torch.Tensor, carry: HierarchicalRea...
    method forward (line 180) | def forward(self, carry: HierarchicalReasoningModel_ACTV1InnerCarry, b...
  class HierarchicalReasoningModel_ACTV1 (line 216) | class HierarchicalReasoningModel_ACTV1(nn.Module):
    method __init__ (line 219) | def __init__(self, config_dict: dict):
    method puzzle_emb (line 225) | def puzzle_emb(self):
    method initial_carry (line 228) | def initial_carry(self, batch: Dict[str, torch.Tensor]):
    method forward (line 240) | def forward(self, carry: HierarchicalReasoningModel_ACTV1Carry, batch:...

FILE: models/layers.py
  function _find_multiple (line 19) | def _find_multiple(a, b):
  function rotate_half (line 23) | def rotate_half(x: torch.Tensor):
  function apply_rotary_pos_emb (line 30) | def apply_rotary_pos_emb(q: torch.Tensor, k: torch.Tensor, cos: torch.Te...
  class CastedLinear (line 43) | class CastedLinear(nn.Module):
    method __init__ (line 44) | def __init__(self,
    method forward (line 58) | def forward(self, input: torch.Tensor) -> torch.Tensor:
  class CastedEmbedding (line 62) | class CastedEmbedding(nn.Module):
    method __init__ (line 63) | def __init__(self,
    method forward (line 76) | def forward(self, input: torch.Tensor) -> torch.Tensor:
  class RotaryEmbedding (line 80) | class RotaryEmbedding(nn.Module):
    method __init__ (line 81) | def __init__(self, dim, max_position_embeddings, base, device=None):
    method forward (line 94) | def forward(self):
  class Attention (line 98) | class Attention(nn.Module):
    method __init__ (line 99) | def __init__(self, hidden_size, head_dim, num_heads, num_key_value_hea...
    method forward (line 112) | def forward(self, cos_sin: CosSin, hidden_states: torch.Tensor) -> tor...
  class SwiGLU (line 138) | class SwiGLU(nn.Module):
    method __init__ (line 139) | def __init__(self, hidden_size: int, expansion: float):
    method forward (line 146) | def forward(self, x):
  function rms_norm (line 151) | def rms_norm(hidden_states: torch.Tensor, variance_epsilon: float) -> to...

FILE: models/losses.py
  function s (line 11) | def s(x, epsilon=1e-30):
  function log_stablemax (line 19) | def log_stablemax(x, dim=-1):
  function stablemax_cross_entropy (line 24) | def stablemax_cross_entropy(logits, labels, ignore_index: int = -100):
  function softmax_cross_entropy (line 34) | def softmax_cross_entropy(logits, labels, ignore_index: int = -100):
  class ACTLossHead (line 40) | class ACTLossHead(nn.Module):
    method __init__ (line 41) | def __init__(self, model: nn.Module, loss_type: str):
    method initial_carry (line 46) | def initial_carry(self, *args, **kwargs):
    method forward (line 49) | def forward(

FILE: models/sparse_embedding.py
  class CastedSparseEmbedding (line 11) | class CastedSparseEmbedding(nn.Module):
    method __init__ (line 12) | def __init__(self, num_embeddings: int, embedding_dim: int, batch_size...
    method forward (line 28) | def forward(self, inputs: torch.Tensor) -> torch.Tensor:
  class CastedSparseEmbeddingSignSGD_Distributed (line 41) | class CastedSparseEmbeddingSignSGD_Distributed(Optimizer):
    method __init__ (line 42) | def __init__(
    method step (line 63) | def step(self, closure=None):  # type: ignore
  function _sparse_emb_signsgd_dist (line 98) | def _sparse_emb_signsgd_dist(

FILE: pretrain.py
  class LossConfig (line 26) | class LossConfig(pydantic.BaseModel):
  class ArchConfig (line 32) | class ArchConfig(pydantic.BaseModel):
  class PretrainConfig (line 39) | class PretrainConfig(pydantic.BaseModel):
  class TrainState (line 74) | class TrainState:
  function create_dataloader (line 84) | def create_dataloader(config: PretrainConfig, split: str, rank: int, wor...
  function create_model (line 108) | def create_model(config: PretrainConfig, train_metadata: PuzzleDatasetMe...
  function cosine_schedule_with_warmup_lr_lambda (line 162) | def cosine_schedule_with_warmup_lr_lambda(
  function init_train_state (line 172) | def init_train_state(config: PretrainConfig, train_metadata: PuzzleDatas...
  function save_train_state (line 190) | def save_train_state(config: PretrainConfig, train_state: TrainState):
  function compute_lr (line 199) | def compute_lr(base_lr: float, config: PretrainConfig, train_state: Trai...
  function train_batch (line 209) | def train_batch(config: PretrainConfig, train_state: TrainState, batch: ...
  function evaluate (line 266) | def evaluate(config: PretrainConfig, train_state: TrainState, eval_loade...
  function save_code_and_config (line 333) | def save_code_and_config(config: PretrainConfig):
  function load_synced_config (line 359) | def load_synced_config(hydra_config: DictConfig, rank: int, world_size: ...
  function launch (line 381) | def launch(hydra_config: DictConfig):

FILE: puzzle_dataset.py
  function _sample_batch (line 14) | def _sample_batch(rng: np.random.Generator, group_order: np.ndarray, puz...
  class PuzzleDatasetConfig (line 41) | class PuzzleDatasetConfig(pydantic.BaseModel):
  class PuzzleDataset (line 53) | class PuzzleDataset(IterableDataset):
    method __init__ (line 54) | def __init__(self, config: PuzzleDatasetConfig, split: str = "train"):
    method _load_metadata (line 68) | def _load_metadata(self) -> PuzzleDatasetMetadata:
    method _lazy_load_dataset (line 72) | def _lazy_load_dataset(self):
    method _collate_batch (line 95) | def _collate_batch(self, batch):
    method _iter_test (line 118) | def _iter_test(self):
    method _iter_train (line 151) | def _iter_train(self):
    method __iter__ (line 189) | def __iter__(self):

FILE: utils/functions.py
  function load_model_class (line 5) | def load_model_class(identifier: str, prefix: str = "models."):
  function get_model_source_path (line 15) | def get_model_source_path(identifier: str, prefix: str = "models."):