diff --git a/models/__init__.py b/models/__init__.py new file mode 100644 index 00000000..7afd4dcf --- /dev/null +++ b/models/__init__.py @@ -0,0 +1,34 @@ +"""Models package for TinyRecursiveModels. + +This package contains the core model implementations including: +- Recursive reasoning models (TRM, HRM, etc.) +- Model layers and utilities +- Loss functions +- Embedding modules + +Classes are available via explicit imports from submodules: + from models.ema import EMAHelper + from models.layers import CastedLinear, Attention + from models.losses import ACTLossHead + from models.sparse_embedding import CastedSparseEmbedding + from models.recursive_reasoning.trm import TinyRecursiveReasoningModel_ACTV1 +""" + +__all__ = [ + # EMA + "EMAHelper", + # Layers + "CastedLinear", + "CastedEmbedding", + "RotaryEmbedding", + "Attention", + "LinearSwish", + "SwiGLU", + "rms_norm", + "CosSin", + # Losses + "ACTLossHead", + # Sparse Embedding + "CastedSparseEmbedding", + "CastedSparseEmbeddingSignSGD_Distributed", +] diff --git a/models/recursive_reasoning/__init__.py b/models/recursive_reasoning/__init__.py new file mode 100644 index 00000000..79aa49c7 --- /dev/null +++ b/models/recursive_reasoning/__init__.py @@ -0,0 +1,51 @@ +"""Recursive reasoning models for TinyRecursiveModels. + +This package contains various recursive reasoning model implementations: +- TRM (Tiny Recursive Model): The core recursive reasoning model +- HRM (Hierarchical Reasoning Model): Alternative hierarchical approach +- Transformer baselines and variants + +Classes are available via explicit imports from submodules: + from models.recursive_reasoning.trm import TinyRecursiveReasoningModel_ACTV1 + from models.recursive_reasoning.hrm import HierarchicalReasoningModel_ACTV1 + from models.recursive_reasoning.transformers_baseline import Model_ACTV2 + from models.recursive_reasoning.trm_ers_pmll import TinyRecursiveReasoningModel_ERS_PMLL +""" + +__all__ = [ + # TRM - Main Tiny Recursive Model + "TinyRecursiveReasoningModel_ACTV1", + "TinyRecursiveReasoningModel_ACTV1_Inner", + "TinyRecursiveReasoningModel_ACTV1Block", + "TinyRecursiveReasoningModel_ACTV1ReasoningModule", + "TinyRecursiveReasoningModel_ACTV1Carry", + "TinyRecursiveReasoningModel_ACTV1InnerCarry", + "TinyRecursiveReasoningModel_ACTV1Config", + + # HRM - Hierarchical Reasoning Model + "HierarchicalReasoningModel_ACTV1", + "HierarchicalReasoningModel_ACTV1_Inner", + "HierarchicalReasoningModel_ACTV1Block", + "HierarchicalReasoningModel_ACTV1ReasoningModule", + "HierarchicalReasoningModel_ACTV1Carry", + "HierarchicalReasoningModel_ACTV1InnerCarry", + "HierarchicalReasoningModel_ACTV1Config", + + # Transformer Baseline + "Model_ACTV2", + "Model_ACTV2_Inner", + "Model_ACTV2Block", + "Model_ACTV2ReasoningModule", + "Model_ACTV2Carry", + "Model_ACTV2InnerCarry", + "Model_ACTV2Config", + + # TRM Variants + "TinyRecursiveReasoningModel_SingleZ", + "TinyRecursiveReasoningModel_Hier6", + + # TRM with ERS and PMLL + "TinyRecursiveReasoningModel_ERS_PMLL", + "MemoryBlock", + "PMLLLatticeState", +]