From b2f4c645635aadb03ff0b0e32d7c993148715f06 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Sat, 31 Jan 2026 04:36:25 +0000
Subject: [PATCH 1/3] Initial plan


From 84580a4e1ae8614c10ea8ea9342765a0ce1d7736 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Sat, 31 Jan 2026 04:40:22 +0000
Subject: [PATCH 2/3] Add package structure with __init__.py files for models

Co-authored-by: drqsatoshi <240532885+drqsatoshi@users.noreply.github.com>
---
 models/__init__.py                     | 34 +++++++++++++++++
 models/recursive_reasoning/__init__.py | 51 ++++++++++++++++++++++++++
 2 files changed, 85 insertions(+)
 create mode 100644 models/__init__.py
 create mode 100644 models/recursive_reasoning/__init__.py

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..fb95b115
--- /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",
+]

From f650fac6ba6654f478388e2177138a0701d0dd75 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Sat, 31 Jan 2026 04:41:28 +0000
Subject: [PATCH 3/3] Fix trailing whitespace in __init__.py

Co-authored-by: drqsatoshi <240532885+drqsatoshi@users.noreply.github.com>
---
 models/recursive_reasoning/__init__.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/models/recursive_reasoning/__init__.py b/models/recursive_reasoning/__init__.py
index fb95b115..79aa49c7 100644
--- a/models/recursive_reasoning/__init__.py
+++ b/models/recursive_reasoning/__init__.py
@@ -21,7 +21,7 @@
     "TinyRecursiveReasoningModel_ACTV1Carry",
     "TinyRecursiveReasoningModel_ACTV1InnerCarry",
     "TinyRecursiveReasoningModel_ACTV1Config",
-    
+
     # HRM - Hierarchical Reasoning Model
     "HierarchicalReasoningModel_ACTV1",
     "HierarchicalReasoningModel_ACTV1_Inner",
@@ -30,7 +30,7 @@
     "HierarchicalReasoningModel_ACTV1Carry",
     "HierarchicalReasoningModel_ACTV1InnerCarry",
     "HierarchicalReasoningModel_ACTV1Config",
-    
+
     # Transformer Baseline
     "Model_ACTV2",
     "Model_ACTV2_Inner",
@@ -39,11 +39,11 @@
     "Model_ACTV2Carry",
     "Model_ACTV2InnerCarry",
     "Model_ACTV2Config",
-    
+
     # TRM Variants
     "TinyRecursiveReasoningModel_SingleZ",
     "TinyRecursiveReasoningModel_Hier6",
-    
+
     # TRM with ERS and PMLL
     "TinyRecursiveReasoningModel_ERS_PMLL",
     "MemoryBlock",