Fix: Bind to 0.0.0.0 and show_api=False to resolve runtime and schema errors

README.md

@@ -1,12 +1,86 @@
 ---
-title: G Ssm Mniah
-emoji: ⚡
-colorFrom: blue
-colorTo: purple
+title: GFN MNIAH Solver
+emoji: 📍
+colorFrom: red
+colorTo: yellow
 sdk: gradio
-sdk_version: 6.9.0
+sdk_version: 4.44.1
+python_version: "3.11"
 app_file: app.py
 pinned: false
+license: cc-by-nc-nd-4.0
+library_name: gfn
+language: en
+pipeline_tag: other
+tags:
+- gfn
+- physics-informed
+- geometric-deep-learning
+- g-ssm
+- needle-in-a-haystack
+- long-context
+model-index:
+- name: gfn-gssm-mniah-k2
+  results:
+  - task:
+      type: other
+      name: Multi-Needle Retrieval
+    dataset:
+      name: synthetic-needle-haystack
+      type: synthetic
+    metrics:
+    - name: Accuracy
+      type: accuracy
+      value: 100.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 📍 GFN MNIAH (Multi-Needle-in-a-Haystack) Solver
+
+[![DOI: 10.5281/zenodo.19141133](https://img.shields.io/badge/DOI-10.5281/zenodo.19141133-blue.svg)](https://doi.org/10.5281/zenodo.19141133)
+[![Models: Hugging Face](https://img.shields.io/badge/Models-Hugging%20Face-orange.svg)](https://huggingface.co/DepthMuun)
+[![GitHub: GFN Framework](https://img.shields.io/badge/GitHub-GFN--Framework-black.svg?logo=github)](https://github.com/DepthMuun/gfn)
+
+This repository demonstrates the **Geometric Flow Network (GFN)** framework's ability to handle extreme-length context retrieval without KV-cache. It implements the **Geodesic State Space Model (G-SSM)** specialized for the K=2 Needle-in-a-Haystack task.
+
+## 🚀 Technical Highlights
+- **Context Length**: Validated up to **32,000 tokens** (linear complexity $O(L)$).
+- **Inductive Bias**: Topologically encodes "needles" as geodetic impulses that curve the world state manifold into a target decision region.
+- **Inference VRAM**: ~38MB constant (O(1) memory).
+
+## 🛠️ Local Installation & Usage
+To run this model locally, you need the **GFN Framework** and the model assets.
+
+### 1. Install GFN Framework
+```bash
+pip install git+https://github.com/DepthMuun/gfn.git
+```
+
+### 2. Clone this repository
+```bash
+git lfs install
+git clone https://huggingface.co/spaces/DepthMuun/gfn-gssm-mniah-k2-space
+cd gfn-gssm-mniah-k2-space
+```
+
+### 3. Run the Interactive Demo
+```bash
+python app.py
+```
+
+## 📜 Citation
+If you use this work, please cite:
+```latex
+@article{sturtz2026geometry,
+  title={Geometric Flow Networks: A Physics-Informed Paradigm for Sequential Intelligence},
+  author={Stürtz, Joaquín},
+  journal={Zenodo Preprints},
+  year={2026},
+  doi={10.5281/zenodo.19141133},
+  url={https://doi.org/10.5281/zenodo.19141133}
+}
+```
+
+## 🔗 Resources
+- **Official Checkpoint**: [GFN MNIAH Model](https://huggingface.co/DepthMuun/gfn-gssm-mniah-k2)
+- **Framework Source**: [GitHub: DepthMuun/gfn](https://github.com/DepthMuun/gfn)
+- **Official Paper**: [Zenodo](https://doi.org/10.5281/zenodo.19141133)

app.py

@@ -0,0 +1,169 @@
+import gradio as gr
+import torch
+import math
+import sys
+import os
+import time
+import json
+from pathlib import Path
+
+# Add local gfn folder to path if it exists (for HF Spaces)
+script_dir = os.path.dirname(os.path.abspath(__file__))
+if os.path.exists(os.path.join(script_dir, "gfn")):
+    sys.path.insert(0, script_dir)
+
+import gfn
+
+def load_model():
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
+    # Load config safely using absolute path
+    config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "config.json")
+    with open(config_path, "r") as f:
+        config = json.load(f)
+        
+    model = gfn.gssm.create(
+        vocab_size=config['architecture']['vocab_size'],
+        dim=config['architecture']['dim'],
+        depth=config['architecture']['depth'],
+        heads=config['architecture']['heads'],
+        integrator=config['architecture']['integrator'],
+        impulse_scale=config['architecture']['impulse_scale'],
+        dynamics_type=config['architecture']['dynamics_type'],
+        topology_type=config['architecture']['topology_type'],
+        physics=config['physics'],
+        holographic=config['architecture'].get('holographic', True),
+    ).to(device)
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    checkpoint_path = os.path.join(script_dir, "mniah_model_final.pt")
+    if not os.path.exists(checkpoint_path):
+        raise FileNotFoundError(f"Missing model weights: {checkpoint_path}. Please place the trained checkpoint here.")
+    
+    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=True)
+    model.load_state_dict(ckpt['model'])
+    model.eval()
+    return model, device
+
+model, device = load_model()
+
+import json
+import tempfile
+
+def run_mniah(seq_len, needle_pos_str, num_needles):
+    try:
+        if needle_pos_str.strip() == "":
+            # Random needles
+            lo = 1
+            pool = torch.randperm(seq_len)[:num_needles]
+            positions = sorted((pool + lo).tolist())
+        else:
+            positions = sorted([int(p.strip()) for p in needle_pos_str.split(",")])
+            if len(positions) != num_needles:
+                return f"Error: Number of positions ({len(positions)}) must match needle count ({num_needles})."
+            if any(p < 1 or p > seq_len for p in positions):
+                return f"Error: Positions must be between 1 and {seq_len}."
+
+        # Input Generation
+        # Format: [Haystack, Needle, ...] (No context token for k2 model)
+        x = torch.zeros(1, seq_len, dtype=torch.long, device=device)
+        
+        # Needles are 1-indexed from the user UI, so subtract 1 for tensor indexing
+        for p in positions:
+            x[0, p - 1] = 1 # Needle token
+            
+        t0 = time.time()
+        with torch.no_grad():
+            output = model(x)
+            x_pred = output[0] # [B, L, D] or [B, L, H, D]
+            if x_pred.ndim == 4:
+                x_pred = x_pred.mean(dim=2)
+                
+        elapsed = time.time() - t0
+        
+        # Binary prediction (toroidal)
+        PI = math.pi
+        TWO_PI = 2.0 * PI
+        half_pi = PI * 0.5
+        dist_pos = torch.min(
+            torch.abs(x_pred - half_pi) % TWO_PI,
+            TWO_PI - (torch.abs(x_pred - half_pi) % TWO_PI)
+        ).mean(dim=-1)
+        dist_neg = torch.min(
+            torch.abs(x_pred + half_pi) % TWO_PI,
+            TWO_PI - (torch.abs(x_pred + half_pi) % TWO_PI)
+        ).mean(dim=-1)
+        
+        preds = (dist_pos < dist_neg).long()[0] # [L]
+        
+        # Result summary
+        last_pos = positions[-1] - 1 # Adjust to 0-indexed
+        acc_after = (preds[last_pos+1:] == 1).float().mean().item() if last_pos+1 < seq_len else 1.0
+        acc_before = (preds[:last_pos] == 0).float().mean().item() if last_pos > 0 else 1.0
+        result_data = {
+            "Status": "Success",
+            "Context Length": seq_len,
+            "Needle Positions": positions,
+            "Predictions (subset)": preds.tolist()[:500],
+            "Acc After Last Needle": f"{acc_after:.2%}",
+            "Acc Before Flip": f"{acc_before:.2%}",
+            "Inference Time": f"{elapsed:.3f}s",
+            "Full Trace": preds.tolist()
+        }
+        
+        # Save to temp file for download
+        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".json", mode='w')
+        json.dump(result_data, temp_file, indent=4)
+        temp_file.close()
+        
+        # Create Markdown table for UI
+        md_summary = f"""
+### 📊 Invariant Evaluation Summary
+
+| Metric | Value |
+|:---|:---|
+| **Status** | 🟢 {result_data['Status']} |
+| **Context Length ($L$)** | {seq_len:,} tokens |
+| **Needles Count ($K$)** | {num_needles} (Positions: {positions}) |
+| **Accuracy (Before Critical Threshold)** | **{acc_before:.2%}** |
+| **Accuracy (After Critical Threshold)** | **{acc_after:.2%}** |
+| **Inference Time** | {elapsed:.3f}s |
+
+> *Note: Model was trained explicitly on $K=2$. Extrapolating to higher $K$ demonstrates how the physical integrator accumulates forces. The 'Accuracy Before' metric measures strict adherence before the final expected needle, but geometrical flips might occur exactly when the integrated energy crosses the learned $K=2$ threshold.*
+"""
+        
+        return md_summary, temp_file.name
+    except Exception as e:
+        return f"### ❌ Error\n{str(e)}", None
+
+with gr.Blocks(title="G-SSM MNIAH Solver") as demo:
+    gr.Markdown("# 📍 G-SSM Multi-Needle-in-a-Haystack (MNIAH)")
+    gr.Markdown("""
+    This space demonstrates the **Geodesic State Space Model (G-SSM)** on the Multi-Needle task.
+    The model must detect and remember exactly **K needles** scattered across a long sequence 
+    of "haystack" tokens. Its internal state only "flips" to the target configuration after **all K impulses** 
+    have been integrated into the geodetic flow.
+    """)
+    
+    with gr.Row():
+        seq_len = gr.Number(value=1000, minimum=64, maximum=10_000_000_000, precision=0, label="Sequence Length (Up to 10B, be mindful of VRAM!)")
+        num_k = gr.Slider(minimum=1, maximum=4, value=2, step=1, label="Number of Needles (K)")
+        positions = gr.Textbox(label="Manual Positions (comma separated, empty for random)", placeholder="e.g. 100, 450")
+    
+    submit_btn = gr.Button("Evaluate Geometric Memory", variant="primary")
+    
+    with gr.Row():
+        with gr.Column(scale=2):
+            output_md = gr.Markdown("### 📊 Results will appear here...")
+        with gr.Column(scale=1):
+            download_btn = gr.File(label="Full Trace (JSON)")
+    
+    submit_btn.click(fn=run_mniah, inputs=[seq_len, positions, num_k], outputs=[output_md, download_btn])
+    
+    gr.Examples(
+        examples=[
+            [1000, "100, 900", 2],
+            [4000, "50, 1500, 3900", 3],
+            [32000, "", 1]
+        ],
+        inputs=[seq_len, positions, num_k]
+    )

config.json

@@ -0,0 +1,58 @@
+{
+  "architecture": {
+    "vocab_size": 2,
+    "dim": 16,
+    "depth": 2,
+    "heads": 2,
+    "integrator": "leapfrog",
+    "impulse_scale": 80.0,
+    "dynamics_type": "direct",
+    "topology_type": "torus",
+    "holographic": true
+  },
+  "physics": {
+    "embedding": {
+      "type": "functional",
+      "mode": "linear",
+      "coord_dim": 16
+    },
+    "readout": {
+      "type": "implicit",
+      "coord_dim": 16
+    },
+    "active_inference": {
+      "enabled": true,
+      "dynamic_time": {
+        "enabled": true
+      },
+      "reactive_curvature": {
+        "enabled": true,
+        "plasticity": 0.2
+      },
+      "singularities": {
+        "enabled": true,
+        "strength": 20.0,
+        "threshold": 0.8
+      }
+    },
+    "fractal": {
+      "enabled": true,
+      "threshold": 0.5,
+      "alpha": 0.2
+    },
+    "topology": {
+      "type": "torus",
+      "riemannian_type": "low_rank"
+    },
+    "stability": {
+      "enable_trace_normalization": true,
+      "base_dt": 0.4,
+      "velocity_saturation": 15.0,
+      "friction": 2.0,
+      "toroidal_curvature_scale": 0.01
+    },
+    "hysteresis": {
+      "enabled": false
+    }
+  }
+}

gfn/__init__.py

@@ -0,0 +1,30 @@
+"""
+GFN (Geodesic Flow Network) Package
+==================================
+Unified framework for Geodesic State Space Models (G-SSM)
+and Inertial State Networks (ISN).
+
+This package implements the GFN paradigm as a platform for
+physics-informed neural dynamics.
+"""
+
+# ── Realizations ──────────────────────────────────────────────────────────────
+from .realizations import api, gssm, isn
+from .realizations.api import create, load, save
+
+# ── Dynamic Registry
+REALIZATIONS = api.list_available()
+
+# ── Package Metadata ──────────────────────────────────────────────────────────
+__version__ = "2.7.0"
+__author__ = "DepthMuun"
+
+__all__ = [
+    "gssm",
+    "isn",
+    "api",
+    "create",
+    "load",
+    "save",
+    "REALIZATIONS",
+]

gfn/realizations/__init__.py

@@ -0,0 +1,19 @@
+"""
+GFN Realizations Subpackage
+===========================
+Contains specific implementations of the GFN paradigm:
+- G-SSM: Geodesic State Space Model (Riemannian/Symplectic)
+- ISN: Inertial State Network (Physics-Informed Interaction Engine)
+"""
+
+from . import api
+from .api import create, list_available
+
+# Trigger registration of standard realizations
+from . import gssm
+from . import isn
+
+# Future realizations can be added here or via external plugins
+# from . import rt
+
+__all__ = ['gssm', 'isn', 'api', 'create', 'list_available']

gfn/realizations/api.py

@@ -0,0 +1,66 @@
+"""
+GFN Realizations API Router (Purified Version)
+==============================================
+Agnostic factory and dynamic registry for GFN architectural realizations.
+Follows SOLID principles: open for extension, closed for modification.
+"""
+
+import logging
+from typing import List, Dict, Any, Optional, Protocol, runtime_checkable
+import torch.nn as nn
+
+logger = logging.getLogger(__name__)
+
+@runtime_checkable
+class RealizationProvider(Protocol):
+    """Protocol defining the interface any GFN realization must provide."""
+    def create(self, **kwargs) -> nn.Module: ...
+    def save(self, model: nn.Module, path: str): ...
+    def load(self, path: str, **kwargs) -> nn.Module: ...
+
+# The Dynamic Registry
+_REGISTRY: Dict[str, RealizationProvider] = {}
+
+def register(name: str, provider: RealizationProvider):
+    """
+    Register a new realization architecture.
+    
+    Args:
+        name: Unique identifier for the realization.
+        provider: An object or module implementing the RealizationProvider protocol.
+    """
+    name = name.lower()
+    if name in _REGISTRY:
+        logger.debug(f"Overwriting GFN realization provider: {name}")
+    _REGISTRY[name] = provider
+
+def list_available() -> List[str]:
+    """List all dynamically registered architectural realizations."""
+    return list(_REGISTRY.keys())
+
+def create(name: str, **kwargs) -> nn.Module:
+    """
+    Unified factory to create any registered GFN realization by name.
+    """
+    name = name.lower()
+    if name not in _REGISTRY:
+        raise ValueError(
+            f"GFN Error: Realization '{name}' is not registered. "
+            f"Ensure the subpackage is imported. Available: {list_available()}"
+        )
+    return _REGISTRY[name].create(**kwargs)
+
+def save(model: nn.Module, path: str, realization: Optional[str] = None):
+    """Unified save interface delegate."""
+    if realization and realization.lower() in _REGISTRY:
+        _REGISTRY[realization.lower()].save(model, path)
+    else:
+        import torch
+        torch.save(model.state_dict(), path)
+
+def load(path: str, realization: str, **kwargs) -> nn.Module:
+    """Unified load interface delegate."""
+    realization = realization.lower()
+    if realization not in _REGISTRY:
+        raise ValueError(f"GFN Error: Realization provider for '{realization}' not found.")
+    return _REGISTRY[realization].load(path, **kwargs)

gfn/realizations/gssm/__init__.py

@@ -0,0 +1,10 @@
+# Export core API
+from .api import create, save, load, Model, Manifold, loss, Trainer
+
+# Register with central realization registry
+try:
+    from .. import api as central_api
+    from . import api as gssm_api
+    central_api.register('gssm', gssm_api)
+except ImportError:
+    pass # Fallback for standalone GSSM usage

gfn/realizations/gssm/api.py

@@ -0,0 +1,102 @@
+"""
+gfn/api.py — GFN V5
+Interfaz pública simplificada y orquestación de alto nivel.
+Centraliza la creación, carga y evaluación de modelos.
+"""
+
+import torch
+import torch.nn as nn
+from typing import Optional, Dict, Any, Union
+
+from .models.factory import ModelFactory
+from .models.manifold import ManifoldModel
+from .losses.factory import LossFactory
+from .training.trainer import GFNTrainer
+from .training.evaluation import ManifoldMetricEvaluator
+
+# -- Alias principales
+Model = ManifoldModel
+Manifold = ManifoldModel
+Trainer = GFNTrainer
+
+def create(*args, **kwargs):
+    """Factory para modelos Manifold (V5)."""
+    return ModelFactory.create(*args, **kwargs)
+
+def loss(config, **kwargs):
+    """Factory para funciones de pérdida (V5)."""
+    return LossFactory.create(config, **kwargs)
+
+def save(model: nn.Module, path: str):
+    """
+    Guarda el modelo y su configuración (HuggingFace Style).
+    """
+    if hasattr(model, 'save_pretrained'):
+        model.save_pretrained(path)
+    else:
+        # Fallback para modelos que no heredan de BaseModel
+        torch.save({'state_dict': model.state_dict()}, path)
+
+def load(path: str, device: Optional[str] = None):
+    """
+    Carga un modelo guardado junto con su configuración.
+    Soporta directorios (HF Style) o archivos .pth/.bin legados.
+    """
+    import os
+    if os.path.isdir(path):
+        return ModelFactory.from_pretrained(path)
+    
+    # Fallback para archivos aislados legados
+    checkpoint = torch.load(path, map_location=device or 'cpu', weights_only=True)
+    config = checkpoint.get('config')
+    if config is None:
+        raise ValueError(f"No se encontró configuración en el checkpoint {path}. Use directorios HF para carga completa.")
+        
+    model = create(config=config)
+    
+    # Robust state_dict extraction (handles different saving conventions)
+    state_dict = checkpoint.get('state_dict') or checkpoint.get('model') or checkpoint
+    
+    # Filter state_dict against the model's actual parameters
+    model_state = model.state_dict()
+    filtered_state = {k: v for k, v in state_dict.items() if k in model_state}
+    
+    # Log filtered keys for debugging (optional)
+    n_filtered = len(state_dict) - len(filtered_state)
+    if n_filtered > 0:
+        import logging
+        logging.getLogger("gssm.api").info(f"Filtered {n_filtered} unexpected keys from state_dict (legacy or auxiliary data).")
+
+    # Load with strict=False to handle potential missing non-essential parameters
+    model.load_state_dict(filtered_state, strict=False)
+    return model
+
+def benchmark(model: nn.Module, dataloader: torch.utils.data.DataLoader, 
+              device: Optional[str] = None) -> Dict[str, float]:
+    """
+    Ejecuta una evaluación rápida de métricas geométricas y de tarea.
+    """
+    device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
+    model.to(device)
+    model.eval()
+    
+    evaluator = ManifoldMetricEvaluator(model)
+    all_x, all_v, all_y = [], [], []
+    
+    with torch.no_grad():
+        for x, y in dataloader:
+            x, y = x.to(device), y.to(device)
+            logits, (xf, vf), info = model(x)
+            
+            all_x.append(xf.detach().cpu())
+            all_v.append(vf.detach().cpu())
+            all_y.append(y.detach().cpu())
+            
+    if not all_x:
+        return {}
+
+    x_total = torch.cat(all_x, dim=0)
+    v_total = torch.cat(all_v, dim=0)
+    y_total = torch.cat(all_y, dim=0)
+    
+    return evaluator.full_report(x_total, v_total, y_total)

gfn/realizations/gssm/config/__init__.py

@@ -0,0 +1,63 @@
+"""
+gfn/config/__init__.py
+Public API for the configuration module — GFN V5
+Centralized configuration system for all GFN components.
+"""
+
+from .schema import (
+    TopologyConfig,
+    StabilityConfig,
+    DynamicTimeConfig,
+    HysteresisConfig,
+    ActiveInferenceConfig,
+    EmbeddingConfig,
+    ReadoutConfig,
+    MixtureConfig,
+    DynamicsConfig,
+    FractalConfig,
+    SingularityConfig,
+    PhysicsConfig,
+    TrainerConfig,
+    ManifoldConfig,
+)
+
+from .defaults import (
+    PHYSICS_DEFAULTS,
+    MODEL_DEFAULTS,
+    TRAINING_DEFAULTS,
+    LOSS_DEFAULTS,
+    get_default,
+)
+
+from .loader import dict_to_physics_config
+from .validator import ConfigValidator, validate_manifold_config, validate_and_print
+
+__all__ = [
+    # Schema classes
+    "TopologyConfig",
+    "StabilityConfig", 
+    "DynamicTimeConfig",
+    "HysteresisConfig",
+    "ActiveInferenceConfig",
+    "EmbeddingConfig",
+    "ReadoutConfig",
+    "MixtureConfig",
+    "DynamicsConfig",
+    "FractalConfig",
+    "SingularityConfig",
+    "PhysicsConfig",
+    "TrainerConfig",
+    "ManifoldConfig",
+    # Defaults
+    "PHYSICS_DEFAULTS",
+    "MODEL_DEFAULTS",
+    "TRAINING_DEFAULTS",
+    "LOSS_DEFAULTS",
+    "get_default",
+    # Loader
+    "dict_to_physics_config",
+    # Validator
+    "ConfigValidator",
+    "validate_manifold_config",
+    "validate_and_print",
+]

gfn/realizations/gssm/config/defaults.py

@@ -0,0 +1,111 @@
+"""
+config/defaults.py — GFN V5
+Valores por defecto centralizados para todas las configuraciones.
+Elimina hardcodes dispersos en implementaciones.
+"""
+
+from typing import Dict, Any
+from ..constants import (
+    DEFAULT_DT, DEFAULT_FRICTION, DEFAULT_PLASTICITY,
+    MAX_VELOCITY, CURVATURE_CLAMP, VELOCITY_FRICTION_SCALE,
+    SINGULARITY_THRESHOLD, BLACK_HOLE_STRENGTH, EPSILON_STANDARD, 
+    TOPOLOGY_TORUS, TOPOLOGY_EUCLIDEAN
+)
+
+# ─── Física ─────────────────────────────────────────────────────────────────
+PHYSICS_DEFAULTS: Dict[str, Any] = {
+    # Topología
+    'topology_type': TOPOLOGY_EUCLIDEAN,
+    'riemannian_type': 'low_rank',
+    'major_radius_R': 2.0,
+    'minor_radius_r': 1.0,
+
+    # Estabilidad — referencias a constants.py (una sola fuente de verdad)
+    'base_dt': DEFAULT_DT,
+    'adaptive_dt': True,
+    'friction': DEFAULT_FRICTION,
+    'velocity_clamp': MAX_VELOCITY,
+    'curvature_clamp': CURVATURE_CLAMP,
+    'enable_trace_normalization': True,
+    'velocity_friction_scale': VELOCITY_FRICTION_SCALE,
+    'integrator_type': 'leapfrog',
+    'friction_mode': 'static',
+
+    # Inferencia activa
+    'active_inference_enabled': True,
+    'holographic_geometry': False,
+    'plasticity': DEFAULT_PLASTICITY,
+
+    # Singularities
+    'singularity_enabled': False,
+    'singularity_threshold': SINGULARITY_THRESHOLD,
+    'singularity_strength': BLACK_HOLE_STRENGTH,
+    'singularity_epsilon': EPSILON_STANDARD,
+
+    # Hysteresis
+    'hysteresis_enabled': False,
+    'hysteresis_decay': 0.95,
+    'hysteresis_ghost_force': True,
+
+    # Stochasticity / Curiosity
+    'stochasticity_enabled': False,
+    'stochasticity_type': 'brownian',
+    'stochasticity_sigma': 0.01,
+    'curiosity_enabled': False,
+    'curiosity_strength': 0.1,
+}
+
+# ─── Modelo ──────────────────────────────────────────────────────────────────
+MODEL_DEFAULTS: Dict[str, Any] = {
+    'dim': 64,
+    'heads': 4,
+    'depth': 2,
+    'rank': 16,
+    'vocab_size': 256,
+    'holographic': False,
+    'pooling_type': None,
+    'initial_spread': 1e-3,
+    'n_trajectories': 1,
+}
+
+# ─── Entrenamiento ───────────────────────────────────────────────────────────
+TRAINING_DEFAULTS: Dict[str, Any] = {
+    'lr': 1e-3,
+    'optimizer_type': 'adam',
+    'weight_decay': 0.0,
+    'grad_clip': 1.0,
+    'epochs': 10,
+    'batch_size': 32,
+    'scheduler_type': 'cosine_warmup',
+    'warmup_steps': 100,
+    'min_lr': 1e-6,
+    'task': 'lm',
+}
+
+# ─── Pérdidas ────────────────────────────────────────────────────────────────
+LOSS_DEFAULTS: Dict[str, Any] = {
+    'type': 'generative',
+    'mode': 'nll',
+    'entropy_coef': 0.0,
+    'label_smoothing': 0.0,
+
+    # Physics-informed
+    'lambda_physics': 0.01,
+    'lambda_geo': 0.001,
+    'lambda_ham': 0.0,
+    'lambda_kin': 0.0,
+}
+
+
+def get_default(section: str, key: str, fallback=None):
+    """
+    Obtiene un valor por defecto desde la sección correspondiente.
+    Uso: get_default('physics', 'base_dt') -> 0.1
+    """
+    mapping = {
+        'physics': PHYSICS_DEFAULTS,
+        'model': MODEL_DEFAULTS,
+        'training': TRAINING_DEFAULTS,
+        'loss': LOSS_DEFAULTS,
+    }
+    return mapping.get(section, {}).get(key, fallback)

gfn/realizations/gssm/config/loader.py

@@ -0,0 +1,163 @@
+"""
+config/loader.py — GFN V5
+Conversión de dicts de configuración a PhysicsConfig tipado.
+Soporte para overrides anidados sobre configs existentes.
+"""
+from typing import Dict, Any, Optional
+from .schema import (
+    PhysicsConfig, TopologyConfig, StabilityConfig, DynamicsConfig,
+    ActiveInferenceConfig, DynamicTimeConfig, HysteresisConfig,
+    EmbeddingConfig, FractalConfig, SingularityConfig,
+)
+
+
+def dict_to_physics_config(d: Dict[str, Any]) -> PhysicsConfig:
+    """
+    Convierte un dict anidado en un PhysicsConfig tipado.
+
+    Soporta todos los sub-campos de PhysicsConfig. Los campos no presentes
+    en el dict mantienen sus valores default del schema.
+    Si `d` ya es PhysicsConfig, lo devuelve intacto.
+    """
+    if isinstance(d, PhysicsConfig):
+        return d
+
+    cfg = PhysicsConfig()
+    _apply_dict_to_physics_config(cfg, d)
+    return cfg
+
+
+def apply_physics_overrides(cfg: PhysicsConfig, overrides: Dict[str, Any]) -> PhysicsConfig:
+    """
+    Aplica un dict de overrides sobre un PhysicsConfig EXISTENTE (in-place).
+
+    A diferencia de dict_to_physics_config(), esta función NO parte de defaults
+    sino que modifica solo los campos presentes en el dict, dejando el resto intacto.
+    Es la función que usa ModelFactory cuando se combina preset + physics kwarg.
+
+    Args:
+        cfg:       PhysicsConfig existente (ej. resultado de get_preset())
+        overrides: Dict anidado con los campos a sobreescribir
+
+    Returns:
+        El mismo cfg modificado in-place (también retornado para encadenamiento).
+    """
+    if not overrides:
+        return cfg
+    _apply_dict_to_physics_config(cfg, overrides)
+    return cfg
+
+
+def _apply_dict_to_physics_config(cfg: PhysicsConfig, d: Dict[str, Any]) -> None:
+    """Función interna — aplica los campos del dict sobre cfg in-place."""
+
+    # ── Topology ──────────────────────────────────────────────────────────────
+    t_d = d.get('topology', d.get('topology_config', {}))
+    if isinstance(t_d, dict) and t_d:
+        _apply(cfg.topology, t_d, [
+            'type', 'R', 'r', 'curvature',
+            'riemannian_type', 'riemannian_rank', 'riemannian_class',
+            'geometry_scope'
+        ])
+        if 'major_radius' in t_d: cfg.topology.R = t_d['major_radius']
+        if 'minor_radius' in t_d: cfg.topology.r = t_d['minor_radius']
+
+    # ── Stability ─────────────────────────────────────────────────────────────
+    s_d = d.get('stability', d.get('stability_config', {}))
+    if isinstance(s_d, dict) and s_d:
+        _apply(cfg.stability, s_d, [
+            'base_dt', 'adaptive', 'dt_min', 'dt_max',
+            'enable_trace_normalization', 'wrap_x',
+            'friction', 'velocity_friction_scale',
+            'curvature_clamp', 'friction_mode',
+            'integrator_type',
+            # alias legacy
+            'velocity_saturation',   # → ignorado, no existe en StabilityConfig
+        ])
+        # Alias de nombres legacy
+        if 'toroidal_curvature_scale' in s_d:
+            cfg.stability.curvature_clamp = s_d['toroidal_curvature_scale']
+
+    # ── Dynamics ──────────────────────────────────────────────────────────────
+    dyn_d = d.get('dynamics', d.get('dynamics_config', {}))
+    if isinstance(dyn_d, dict) and dyn_d:
+        if 'type' in dyn_d:
+            cfg.dynamics.type = dyn_d['type']
+
+    # ── Active Inference ──────────────────────────────────────────────────────
+    ai_d = d.get('active_inference', d.get('active_inference_config', {}))
+    if isinstance(ai_d, dict) and ai_d:
+        _apply(cfg.active_inference, ai_d, [
+            'enabled', 'holographic_geometry',
+            'thermodynamic_geometry', 'plasticity',
+        ])
+        # Dynamic time
+        dt_d = ai_d.get('dynamic_time', {})
+        if isinstance(dt_d, dict) and dt_d:
+            _apply(cfg.active_inference.dynamic_time, dt_d, ['enabled', 'type'])
+        # Reactive curvature — es un dict interno
+        rc_d = ai_d.get('reactive_curvature', {})
+        if isinstance(rc_d, dict) and rc_d:
+            cfg.active_inference.reactive_curvature.update(rc_d)
+        # Stochasticity — es un dict interno
+        st_d = ai_d.get('stochasticity', {})
+        if isinstance(st_d, dict) and st_d:
+            cfg.active_inference.stochasticity.update(st_d)
+        # Curiosity — es un dict interno
+        cu_d = ai_d.get('curiosity', {})
+        if isinstance(cu_d, dict) and cu_d:
+            cfg.active_inference.curiosity.update(cu_d)
+    # ── Hysteresis (pueden estar en raíz O dentro de active_inference) ────────
+    hyst_src = d.get('hysteresis', ai_d.get('hysteresis', {}) if isinstance(ai_d, dict) else {})
+    if isinstance(hyst_src, dict) and hyst_src:
+        _apply(cfg.hysteresis, hyst_src, [
+            'enabled', 'ghost_force', 'hyst_decay',
+            'hyst_update_w', 'hyst_update_b',
+            'hyst_readout_w', 'hyst_readout_b',
+        ])
+
+    # ── Singularities (pueden estar en raíz O dentro de active_inference) ─────
+    sing_src = d.get('singularities', ai_d.get('singularities', {}) if isinstance(ai_d, dict) else {})
+    if isinstance(sing_src, dict) and sing_src:
+        _apply(cfg.singularities, sing_src, [
+            'enabled', 'epsilon', 'strength', 'threshold'
+        ])
+
+    # ── Embedding ─────────────────────────────────────────────────────────────
+    emb_d = d.get('embedding', d.get('embedding_config', {}))
+    if isinstance(emb_d, dict) and emb_d:
+        _apply(cfg.embedding, emb_d, [
+            'type', 'mode', 'coord_dim', 'impulse_scale', 'omega_0'
+        ])
+
+    # ── Readout ───────────────────────────────────────────────────────────────
+    read_d = d.get('readout', d.get('readout_config', {}))
+    if isinstance(read_d, dict) and read_d:
+        _apply(cfg.readout, read_d, ['type'])
+
+    # ── Mixture ───────────────────────────────────────────────────────────────
+    mix_d = d.get('mixture', d.get('mixture_config', {}))
+    if isinstance(mix_d, dict) and mix_d:
+        _apply(cfg.mixture, mix_d, ['coupler_mode'])
+
+    # ── Fractal ───────────────────────────────────────────────────────────────
+    frac_d = d.get('fractal', {})
+    if isinstance(frac_d, dict) and frac_d:
+        _apply(cfg.fractal, frac_d, ['enabled', 'threshold', 'alpha'])
+
+    # ── Top-level trajectory_mode ─────────────────────────────────────────────
+    if 'trajectory_mode' in d:
+        cfg.trajectory_mode = d['trajectory_mode']
+
+    # ── Attention/mixer alias (legacy ECG configs) ────────────────────────────
+    # 'attention': {'mixer_type': 'low_rank'} — se ignora acá, aplica en ManifoldConfig
+
+
+def _apply(target, source: dict, keys: list) -> None:
+    """Copia las claves presentes en source hacia target (setattr)."""
+    for k in keys:
+        if k in source:
+            try:
+                setattr(target, k, source[k])
+            except AttributeError:
+                pass  # clave no existe en el dataclass — ignorar silenciosamente

gfn/realizations/gssm/config/schema.py

@@ -0,0 +1,199 @@
+# schema.py — GFN V5
+# Definiciones de clases de configuración (Schema)
+# SEPARACIÓN: Los valores por defecto van a defaults.py, las constantes físicas a constants.py
+
+from dataclasses import dataclass, field, asdict
+from typing import Dict, Any, Optional, List
+
+# Importar constantes físicas正确adas
+from ..constants import (
+    EPSILON_STANDARD,
+    TOPOLOGY_TORUS,
+    MIN_DT,
+    MAX_DT,
+    CURVATURE_CLAMP,
+    SINGULARITY_THRESHOLD,
+    BLACK_HOLE_STRENGTH,
+    DEFAULT_DT,
+    DEFAULT_FRICTION,
+    DEFAULT_PLASTICITY,
+    MAX_VELOCITY,
+)
+
+
+@dataclass
+class TopologyConfig:
+    """Configuración de topología del manifold."""
+    type: str = TOPOLOGY_TORUS 
+    R: float = 2.0           # Radio mayor del toro (default)
+    r: float = 1.0           # Radio menor del toro (default)
+    curvature: float = 0.0
+    riemannian_type: str = 'reactive' 
+    riemannian_rank: int = 16
+    riemannian_class: Optional[str] = None
+    geometry_scope: str = 'local'  # 'local' (dim/heads) or 'global' (full dim)
+    # NUEVO: Parámetros aprendibles
+    learnable_R: bool = True   # Hacer R aprendible (como dice el paper)
+    learnable_r: bool = True   # Hacer r aprendible (como dice el paper)
+
+
+@dataclass
+class StabilityConfig:
+    """Configuración de estabilidad numérica."""
+    base_dt: float = DEFAULT_DT 
+    adaptive: bool = True
+    dt_min: float = MIN_DT
+    dt_max: float = MAX_DT
+    enable_trace_normalization: bool = True
+    wrap_x: bool = True
+    friction: float = DEFAULT_FRICTION 
+    velocity_friction_scale: float = 0.0
+    velocity_saturation: float = 0.0  # P2.3: 0 = disabled, >0 = clamp magnitude via tanh
+    curvature_clamp: float = CURVATURE_CLAMP
+    friction_mode: str = 'static'  # 'static' or 'lif'
+    integrator_type: str = 'leapfrog'
+    toroidal_curvature_scale: float = 0.01  # scale for torus Christoffel contribution
+
+
+@dataclass
+class DynamicTimeConfig:
+    enabled: bool = False
+    type: str = 'riemannian'
+
+
+@dataclass
+class HysteresisConfig:
+    enabled: bool = False
+    ghost_force: bool = True
+    hyst_decay: float = 0.1
+    hyst_update_w: float = 1.0
+    hyst_update_b: float = 0.0
+    hyst_readout_w: float = 1.0
+    hyst_readout_b: float = 0.0
+
+
+@dataclass
+class ActiveInferenceConfig:
+    enabled: bool = False
+    holographic_geometry: bool = False 
+    thermodynamic_geometry: bool = False
+    plasticity: float = DEFAULT_PLASTICITY
+    dynamic_time: DynamicTimeConfig = field(default_factory=DynamicTimeConfig)
+    reactive_curvature: Dict[str, Any] = field(default_factory=lambda: {
+        "enabled": False, 
+        "plasticity": 0.0
+    })
+    geodesic_lensing: Dict[str, Any] = field(default_factory=lambda: {"enabled": False})
+    
+    # Exploration / Noise
+    stochasticity: Dict[str, Any] = field(default_factory=lambda: {
+        "enabled": False, 
+        "type": "brownian", 
+        "sigma": 0.01, 
+        "theta": 0.15, 
+        "mu": 0.0
+    })
+    curiosity: Dict[str, Any] = field(default_factory=lambda: {
+        "enabled": False, 
+        "strength": 0.1, 
+        "decay": 0.99
+    })
+
+
+@dataclass
+class EmbeddingConfig:
+    type: str = 'standard'
+    mode: str = 'linear'
+    coord_dim: int = 16
+    impulse_scale: float = 1.0
+    omega_0: float = 30.0
+
+
+@dataclass
+class ReadoutConfig:
+    type: str = 'standard'
+
+
+@dataclass
+class MixtureConfig:
+    coupler_mode: str = 'mean_field'
+
+
+@dataclass
+class DynamicsConfig:
+    type: str = 'direct'
+
+
+@dataclass
+class FractalConfig:
+    enabled: bool = False
+    threshold: float = 0.5
+    alpha: float = 0.2
+
+
+@dataclass
+class SingularityConfig:
+    enabled: bool = False
+    epsilon: float = EPSILON_STANDARD
+    strength: float = BLACK_HOLE_STRENGTH
+    threshold: float = SINGULARITY_THRESHOLD
+
+
+@dataclass
+class PhysicsConfig:
+    """Configuración completa de física."""
+    topology: TopologyConfig = field(default_factory=TopologyConfig)
+    stability: StabilityConfig = field(default_factory=StabilityConfig)
+    dynamics: DynamicsConfig = field(default_factory=DynamicsConfig)
+    active_inference: ActiveInferenceConfig = field(default_factory=ActiveInferenceConfig)
+    embedding: EmbeddingConfig = field(default_factory=EmbeddingConfig)
+    readout: ReadoutConfig = field(default_factory=ReadoutConfig)
+    mixture: MixtureConfig = field(default_factory=MixtureConfig)
+    fractal: FractalConfig = field(default_factory=FractalConfig)
+    hysteresis: HysteresisConfig = field(default_factory=HysteresisConfig)
+    singularities: SingularityConfig = field(default_factory=SingularityConfig)
+    trajectory_mode: str = 'partition'
+    lensing: Dict[str, Any] = field(default_factory=lambda: {'enabled': False})
+    checkpointing: Dict[str, Any] = field(default_factory=lambda: {'enabled': False})
+
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)
+
+
+@dataclass
+class TrainerConfig:
+    lr: float = 1e-3
+    optimizer: str = 'adamw'
+    max_lr: Optional[float] = None
+    total_steps: Optional[int] = None
+    loss_config: Dict[str, Any] = field(default_factory=lambda: {
+        'lambda_g': 0.001,
+        'lambda_h': 0.0,
+        'geodesic_mode': 'magnitude'
+    })
+
+
+@dataclass
+class ManifoldConfig:
+    """Configuración principal del modelo Manifold."""
+    vocab_size: int
+    dim: int = 512
+    depth: int = 4
+    heads: int = 4
+    rank: int = 32
+    integrator: str = 'leapfrog'
+    physics: PhysicsConfig = field(default_factory=PhysicsConfig)
+    trainer: TrainerConfig = field(default_factory=TrainerConfig)
+    adjoint_enabled: bool = False
+    adjoint_rtol: float = 1e-4
+    adjoint_atol: float = 1e-4
+    holographic: bool = False
+    impulse_scale: float = 1.0
+    dynamics_type: str = 'direct'
+    mixer_type: str = 'low_rank'
+    trajectory_mode: str = 'partition'
+    coupler_mode: str = 'mean_field'
+    initial_spread: float = 1e-3
+
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)

gfn/realizations/gssm/config/serialization.py

@@ -0,0 +1,39 @@
+import dataclasses
+from typing import Any, Dict, Type, TypeVar, get_type_hints, get_args, get_origin, Union
+
+T = TypeVar('T')
+
+def from_dict(cls: Type[T], data: Dict[str, Any]) -> T:
+    """
+    Reconstructs a nested dataclass from a dictionary.
+    Handles nested dataclasses and basic types.
+    """
+    if not dataclasses.is_dataclass(cls):
+        return data
+
+    field_types = get_type_hints(cls)
+    kwargs = {}
+    
+    for field in dataclasses.fields(cls):
+        if field.name in data:
+            value = data[field.name]
+            field_type = field_types[field.name]
+            
+            # Handle Optional[T]
+            origin = get_origin(field_type)
+            if origin is Union:
+                args = get_args(field_type)
+                if type(None) in args:
+                    # It's an Optional. Find the non-None type
+                    field_type = [arg for arg in args if arg is not type(None)][0]
+            
+            # Handle nested dataclasses
+            if dataclasses.is_dataclass(field_type):
+                if value is not None:
+                    kwargs[field.name] = from_dict(field_type, value)
+                else:
+                    kwargs[field.name] = None
+            else:
+                kwargs[field.name] = value
+                
+    return cls(**kwargs)

gfn/realizations/gssm/config/validator.py

@@ -0,0 +1,109 @@
+"""
+Validación de configuraciones — GFN V5
+Verifica la compatibilidad de parámetros antes de construir componentes.
+Fusionado de utils/validation.py y config/validator.py original.
+"""
+
+from typing import Dict, Any, List, Optional
+from .schema import ManifoldConfig, PhysicsConfig
+from ..constants import TOPOLOGY_TORUS, TOPOLOGY_EUCLIDEAN, TOPOLOGY_SPHERE
+
+class ConfigValidationError(Exception):
+    """Error de validación de configuración crítica."""
+    pass
+
+class ConfigValidator:
+    """Central validator for GFN configurations."""
+    
+    @staticmethod
+    def validate_physics(cfg: PhysicsConfig, dim: Optional[int] = None, heads: Optional[int] = None):
+        """
+        Validate physical and architectural consistency of PhysicsConfig.
+        Raises ConfigValidationError if strict topology/stability rules are violated.
+        """
+        # 1. Topology checks
+        if cfg.topology.type == TOPOLOGY_TORUS:
+            if dim is not None and heads is not None:
+                head_dim = dim // heads
+                if head_dim % 2 != 0:
+                    raise ConfigValidationError(
+                        f"Toroid geometry requires head_dim (dim//heads) to be even. "
+                        f"Found {dim}//{heads}={head_dim}"
+                    )
+            
+        if cfg.topology.type == TOPOLOGY_SPHERE and cfg.topology.curvature <= 0:
+             raise ConfigValidationError("Spherical topology requires positive curvature.")
+
+        # 2. Stability checks
+        if cfg.stability.base_dt <= 0:
+            raise ConfigValidationError("base_dt must be positive.")
+        if cfg.stability.friction < 0:
+            raise ConfigValidationError("friction cannot be negative.")
+        
+        # 3. Mode Compatibility
+        if cfg.trajectory_mode == 'ensemble' and heads is not None and heads <= 1:
+            raise ConfigValidationError("Ensemble trajectory mode requires more than 1 head.")
+
+
+def validate_manifold_config(config: ManifoldConfig) -> List[str]:
+    """
+    Valida un ManifoldConfig completo y su PhysicsConfig anidado.
+    Retorna lista de warnings (vacía si todo está OK).
+    Lanza ConfigValidationError en errores críticos o de compatibilidad.
+    """
+    warnings = []
+
+    # Validaciones críticas (Raise exceptions)
+    if config.dim % config.heads != 0:
+        raise ConfigValidationError(
+            f"dim={config.dim} no es divisible por heads={config.heads}. "
+            f"head_dim={config.dim/config.heads:.1f} no es entero."
+        )
+
+    if config.vocab_size <= 0:
+        raise ConfigValidationError(f"vocab_size={config.vocab_size} debe ser > 0.")
+        
+    if config.depth <= 0:
+        raise ConfigValidationError(f"depth={config.depth} debe ser > 0.")
+
+    # Validate Physics properties via centralized method
+    ConfigValidator.validate_physics(config.physics, config.dim, config.heads)
+
+    # Validaciones suaves (Warnings)
+    head_dim = config.dim // config.heads
+    topo_type = config.physics.topology.type.lower()
+    
+    if topo_type == TOPOLOGY_TORUS and head_dim % 2 != 0:
+        warnings.append(
+            f"[WARN] Para geometría toroidal, head_dim={head_dim} debería ser par "
+            f"para representaciones sin/cos. Considera usar heads={config.dim // (head_dim + 1)} o similar."
+        )
+
+    if config.rank > config.dim:
+        warnings.append(
+            f"[WARN] rank={config.rank} > dim={config.dim}. "
+            f"La descomposición no es de rango bajo. ¿Intencional?"
+        )
+
+    dt = config.physics.stability.base_dt
+    if dt > 1.0:
+        warnings.append(f"[WARN] base_dt={dt} > 1.0 puede causar inestabilidad numérica.")
+    if dt < 1e-5:
+        warnings.append(f"[WARN] base_dt={dt} < 1e-5 puede ralentizar convergencia.")
+
+    return warnings
+
+
+def validate_and_print(config: ManifoldConfig) -> bool:
+    """
+    Valida la configuración e imprime warnings.
+    Retorna True si es válida, False si hubo errores.
+    """
+    try:
+        warnings = validate_manifold_config(config)
+        for w in warnings:
+            print(w)
+        return True
+    except ConfigValidationError as e:
+        print(f"[CONFIG ERROR] {e}")
+        return False

gfn/realizations/gssm/constants.py

@@ -0,0 +1,67 @@
+# constants.py — GFN V5
+# Constantes físicas y matemáticas universales.
+# NO contiene hiperparámetros de entrenamiento ni valores configurables.
+
+import torch
+
+# ─── Constantes Matemáticas ─────────────────────────────────────────────────
+PI = 3.14159265358979
+E = 2.718281828459045
+SQRT_2 = 1.4142135623730951
+LOG_2 = 0.6931471805599453
+
+# ─── Estabilidad Numérica ─────────────────────────────────────────────────
+EPS = 1e-8
+INF = 1e12
+EPSILON_STANDARD = 1e-7
+EPSILON_SMOOTH = 1e-9
+EPSILON_STRONG = 1e-6
+CLAMP_MIN_STRONG = 1e-4
+
+# ─── Límites Físicos ───────────────────────────────────────────────────────
+MIN_DT = 0.001
+MAX_DT = 1.0
+
+# ─── Geometría / Curvatura ─────────────────────────────────────────────────
+CURVATURE_CLAMP = 5.0  # Maximum absolute value of Christoffel output
+FRICTION_SCALE = 0.1   # Global friction scaling factor
+VELOCITY_FRICTION_SCALE = 0.01
+
+# ─── Gate initialization constants ─────────────────────────────────────────
+GATE_BIAS_OPEN   = 2.0   # sigmoid(2.0)  ≈ 0.88
+GATE_BIAS_CLOSED = -2.0  # sigmoid(-2.0) ≈ 0.12
+
+# ─── Singularity / Active Inference ───────────────────────────────────────
+SINGULARITY_THRESHOLD = 0.5
+BLACK_HOLE_STRENGTH = 3.0
+SINGULARITY_GATE_SLOPE = 10.0
+
+# ─── Torus geometry ───────────────────────────────────────────────────────
+TOROIDAL_MAJOR_RADIUS   = 1.0
+TOROIDAL_MINOR_RADIUS   = 0.3
+TOROIDAL_PERIOD         = 2.0 * PI
+TOROIDAL_CURVATURE_SCALE = 0.1
+
+# ─── Tipo de dato por defecto ─────────────────────────────────────────────
+DTYPE = torch.float32
+
+# ─── Topology Names ───────────────────────────────────────────────────────
+TOPOLOGY_TORUS      = "torus"
+TOPOLOGY_SPHERE     = "spherical"
+TOPOLOGY_HYPERBOLIC = "hyperbolic"
+TOPOLOGY_EUCLIDEAN  = "euclidean"
+
+# ─── Dynamics Modes ───────────────────────────────────────────────────────
+DYNAMICS_DIRECT     = "direct"
+DYNAMICS_RESIDUAL   = "residual"
+DYNAMICS_MIX        = "mix"
+DYNAMICS_GATED      = "gated"
+DYNAMICS_STOCHASTIC = "stochastic"
+
+# ─── Alias de compatibilidad (valores por defecto que moved a defaults.py) ─
+# NOTA: Estos valores se mantienen aquí por compatibilidad pero deberían
+# imports desde config/defaults.py en código nuevo
+DEFAULT_FRICTION = 0.01
+DEFAULT_DT = 0.1
+DEFAULT_PLASTICITY = 0.05
+MAX_VELOCITY = 10.0

gfn/realizations/gssm/core/__init__.py

@@ -0,0 +1,7 @@
+"""
+core/__init__.py — GFN V5
+"""
+from ..core.types import ManifoldState, Trajectory, StepResult, ModelOutput
+from ..core.state import ManifoldStateManager
+
+__all__ = ['ManifoldState', 'Trajectory', 'StepResult', 'ModelOutput', 'ManifoldStateManager']

gfn/realizations/gssm/core/state.py

@@ -0,0 +1,60 @@
+"""
+core/state.py — GFN V5
+Manejo de estado del manifold (posición + velocidad).
+"""
+
+import torch
+import torch.nn as nn
+from typing import Optional, Tuple
+
+
+class ManifoldStateManager:
+    """
+    Gestiona la inicialización y manipulación del estado (x, v).
+    Compatible con batches y múltiples cabezales.
+    """
+
+    @staticmethod
+    def initialize(x0: nn.Parameter, v0: nn.Parameter,
+                   batch_size: int, n_trajectories: int = 1,
+                   initial_spread: float = 1e-3) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Inicializa el estado (x, v) para un batch dado.
+
+        Args:
+            x0, v0:         Parámetros iniciales [1, H, HD]
+            batch_size:     Tamaño del batch
+            n_trajectories: Número de trayectorias paralelas
+            initial_spread: Ruido inicial
+
+        Returns:
+            (x, v) — [B, H, HD]
+        """
+        x = x0.expand(batch_size, -1, -1)
+        v = v0.expand(batch_size, -1, -1)
+
+        if initial_spread > 0:
+            x = x + torch.randn_like(x) * initial_spread
+
+        return x.contiguous(), v.contiguous()
+
+    @staticmethod
+    def from_tuple(state: Optional[Tuple], x0: nn.Parameter, v0: nn.Parameter,
+                   batch_size: int, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Construye (x, v) desde un estado previo o desde parámetros iniciales.
+        Compatible con el API de BasicModel.
+        """
+        if state is not None and isinstance(state, (tuple, list)) and len(state) == 2:
+            return state[0], state[1]
+        return ManifoldStateManager.initialize(x0, v0, batch_size, **kwargs)
+
+    @staticmethod
+    def wrap_torus(x: torch.Tensor) -> torch.Tensor:
+        """Proyecta posición al dominio toroidal [-π, π]."""
+        return torch.atan2(torch.sin(x), torch.cos(x))
+
+    @staticmethod
+    def energy(v: torch.Tensor) -> torch.Tensor:
+        """Energía cinética H = 0.5 * ||v||² por muestra."""
+        return 0.5 * (v ** 2).sum(dim=-1)

gfn/realizations/gssm/core/types.py

@@ -0,0 +1,27 @@
+"""
+core/types.py — GFN V5
+Tipos y type aliases del framework.
+"""
+
+import torch
+from typing import Dict, Any, Tuple, Optional, List, Union
+
+# ─── State types ─────────────────────────────────────────────────────────────
+# (position, velocity) pair
+ManifoldState = Tuple[torch.Tensor, torch.Tensor]
+
+# Trajectory: list of (x, v) states over time
+Trajectory = List[ManifoldState]
+
+# Force tensor (same shape as x, v)
+Force = torch.Tensor
+
+# Integration step result
+StepResult = Dict[str, torch.Tensor]  # {'x': ..., 'v': ...}
+
+# ─── Config types ─────────────────────────────────────────────────────────────
+ConfigDict = Dict[str, Any]
+
+# ─── Forward pass outputs ─────────────────────────────────────────────────────
+# (logits, state, info_dict)
+ModelOutput = Tuple[torch.Tensor, ManifoldState, Dict[str, Any]]

gfn/realizations/gssm/csrc/README.md

@@ -0,0 +1,2 @@
+# Native C++/CUDA Extensions
+Store raw .cpp and .cu files here. Bindings remain in new_gfn/cuda/

gfn/realizations/gssm/csrc/compile_cuda_12.9.bat

@@ -0,0 +1,68 @@
+@echo off
+
+REM Ensure we are in the script's directory so we run the correct setup.py
+cd /d "%~dp0"
+
+echo ================================================================
+echo [GFN] Custom CUDA Kernel Compilation Pipeline (VS 2022 + CUDA 12.9)
+echo ================================================================
+
+REM --- 1. Find Visual Studio 2022 Installation ---
+set "VS_PATH="
+
+if exist "C:\Program Files\Microsoft Visual Studio\2022\Community\VC\Auxiliary\Build\vcvars64.bat" (
+    set "VS_PATH=C:\Program Files\Microsoft Visual Studio\2022\Community\VC\Auxiliary\Build\vcvars64.bat"
+) else if exist "C:\Program Files\Microsoft Visual Studio\2022\Professional\VC\Auxiliary\Build\vcvars64.bat" (
+    set "VS_PATH=C:\Program Files\Microsoft Visual Studio\2022\Professional\VC\Auxiliary\Build\vcvars64.bat"
+) else if exist "C:\Program Files\Microsoft Visual Studio\2022\Enterprise\VC\Auxiliary\Build\vcvars64.bat" (
+    set "VS_PATH=C:\Program Files\Microsoft Visual Studio\2022\Enterprise\VC\Auxiliary\Build\vcvars64.bat"
+)
+
+if "%VS_PATH%"=="" (
+    echo [ERROR] Could not find Visual Studio 2022 installation.
+    pause
+    exit /b 1
+)
+
+echo [*] Found MSVC Environment: "%VS_PATH%"
+echo [*] Initializing Developer Console...
+call "%VS_PATH%"
+
+REM --- 2. Setup CUDA Environment (12.9) ---
+set "CUDA_PATH=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.9"
+set "PATH=%CUDA_PATH%\bin;%PATH%"
+
+echo [*] Using CUDA Path: "%CUDA_PATH%"
+nvcc --version
+
+REM --- 3. Compile Kernels ---
+echo [*] Cleaning old builds...
+rmdir /s /q build
+rmdir /s /q gfn_cuda.egg-info
+del /q *.pyc
+rmdir /s /q __pycache__
+
+echo.
+echo [*] Starting Setup Compilation (In-place)...
+
+REM Fix for "It seems that the VC environment is activated..." warning
+set DISTUTILS_USE_SDK=1
+
+
+
+python  setup.py build_ext --inplace 
+
+if %errorlevel% neq 0 (
+    echo [ERROR] Compilation failed.
+    echo Ensure you have PyTorch installed for CUDA 12.x:
+    echo pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124
+    pause
+    exit /b 1
+)
+
+echo.
+echo [SUCCESS] Kernels compiled to local .pyd file!
+echo Verified import:
+
+python -c "print('SUCCESS: gfn_cuda module imported directly!')"
+pause

gfn/realizations/gssm/csrc/extension.cpp

@@ -0,0 +1,87 @@
+#include <torch/extension.h>
+#include <vector>
+#include "integrators/integrators.h"
+
+// Declaración de funciones (Toroidal Loss)
+torch::Tensor toroidal_distance_loss_fwd(const torch::Tensor& y_pred, const torch::Tensor& y_true);
+torch::Tensor toroidal_distance_loss_bwd(const torch::Tensor& grad_output, const torch::Tensor& y_pred, const torch::Tensor& y_true);
+
+// Declaración de funciones (Low Rank Christoffel)
+torch::Tensor low_rank_christoffel_fwd(
+    const torch::Tensor& v, const torch::Tensor& U, const torch::Tensor& W,
+    double clamp_val, bool enable_trace_norm, bool is_paper_version);
+
+// Implementación de Backward puro en ATen C++ (Compilado por MSVC, evadiendo bug de NVCC CICC)
+std::vector<torch::Tensor> low_rank_christoffel_bwd(
+    const torch::Tensor& grad_gamma, 
+    const torch::Tensor& v, 
+    const torch::Tensor& U, 
+    const torch::Tensor& W,
+    const torch::Tensor& gamma_out, 
+    double clamp_val, 
+    bool enable_trace_norm, 
+    bool is_paper_version) 
+{
+    // Fast pure ATen operations avoiding Python overhead
+    auto g_norm = gamma_out / clamp_val;
+    auto d_tanh = 1.0 - g_norm.pow(2);
+    auto grad_raw = grad_gamma * d_tanh;  // [B, H, D]
+
+    if (enable_trace_norm) {
+        auto mean_d = grad_raw.mean(-1, /*keepdim=*/true);
+        grad_raw = grad_raw - mean_d;
+    }
+
+    // Explicit Batched Matrix Multiplication (bmm) to avoid matmul broadcast crashes
+    // W is [H, D, R], grad_raw is [B, H, D]
+    auto grad_raw_h = grad_raw.permute({1, 0, 2}); // [H, B, D]
+    auto d_sq_h = torch::bmm(grad_raw_h, W);       // [H, B, D] @ [H, D, R] -> [H, B, R]
+    auto d_sq = d_sq_h.permute({1, 0, 2});         // [B, H, R]
+    
+    auto v_h = v.permute({1, 0, 2});               // [H, B, D]
+    auto v_r_h = torch::bmm(v_h, U);               // [H, B, D] @ [H, D, R] -> [H, B, R]
+    auto v_r = v_r_h.permute({1, 0, 2});           // [B, H, R]
+    
+    torch::Tensor d_vr;
+    if (is_paper_version) {
+        auto vr_norm = torch::norm(v_r, 2, -1, true);
+        auto denom = 1.0 + vr_norm;
+        
+        // Correct Chain Rule for normalized denominator Coupling:
+        // grad_vr_j = grad_phi_j * (2v_j / denom) - v_j * Sum_k(grad_phi_k * v_k^2) / (||v|| * denom^2)
+        auto S = (d_sq * v_r.pow(2)).sum(-1, true);
+        auto term1 = d_sq * (2.0 * v_r / denom);
+        auto term2 = (v_r * S) / (vr_norm * denom.pow(2) + 1e-8);
+        d_vr = term1 - term2;
+    } else {
+        d_vr = d_sq * 2.0 * v_r;
+    }
+
+    auto d_vr_h = d_vr.permute({1, 0, 2});         // [H, B, R]
+    auto U_t = U.transpose(-1, -2);                // [H, R, D]
+    auto d_v_h = torch::bmm(d_vr_h, U_t);          // [H, B, R] @ [H, R, D] -> [H, B, D]
+    auto d_v = d_v_h.permute({1, 0, 2});           // [B, H, D]
+
+    // We accumulate W and U gradients over Batch:
+    auto sq = is_paper_version ? v_r.pow(2) / (1.0 + torch::norm(v_r, 2, -1, true)) : v_r.pow(2); // [B, H, R]
+    auto sq_h = sq.permute({1, 0, 2});             // [H, B, R]
+    
+    auto grad_raw_h_t = grad_raw_h.transpose(-1, -2); // [H, D, B]
+    auto d_W = torch::bmm(grad_raw_h_t, sq_h);        // [H, D, B] @ [H, B, R] -> [H, D, R]
+
+    auto v_h_t = v_h.transpose(-1, -2);               // [H, D, B]
+    auto d_U = torch::bmm(v_h_t, d_vr_h);             // [H, D, B] @ [H, B, R] -> [H, D, R]
+
+    return {d_v, d_U, d_W};
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("toroidal_distance_loss_fwd", &toroidal_distance_loss_fwd, "Toroidal Distance Loss Forward (CUDA)");
+    m.def("toroidal_distance_loss_bwd", &toroidal_distance_loss_bwd, "Toroidal Distance Loss Backward (CUDA)");
+    
+    m.def("low_rank_christoffel_fwd", &low_rank_christoffel_fwd, "Low Rank Christoffel Forward Kernel");
+    m.def("low_rank_christoffel_bwd", &low_rank_christoffel_bwd, "Low Rank Christoffel Backward ATen");
+    
+    m.def("yoshida_fwd", &yoshida_fwd_aten, "Yoshida C++ Macro Integrator Step");
+    m.def("leapfrog_fwd", &leapfrog_fwd_aten, "Leapfrog C++ Macro Integrator Step");
+}

gfn/realizations/gssm/csrc/geometry/low_rank.cu

@@ -0,0 +1,160 @@
+#include <torch/extension.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+template <typename scalar_t>
+__global__ void low_rank_christoffel_fwd_kernel(
+    const scalar_t* __restrict__ v,  // [B, H, D]
+    const scalar_t* __restrict__ U,  // [H, D, R]
+    const scalar_t* __restrict__ W,  // [H, D, R]
+    scalar_t* __restrict__ gamma,    // [B, H, D]
+    const int B, const int H, const int D, const int R,
+    const scalar_t clamp_val,
+    const bool enable_trace_norm,
+    const bool is_paper_version) 
+{
+    // A block computes the output for one (b, h) pair
+    int bh = blockIdx.x;
+    if (bh >= B * H) return;
+    
+    int h = bh % H;
+    
+    // Dynamic shared memory allocations
+    extern __shared__ char smem[];
+    scalar_t* v_s_d     = reinterpret_cast<scalar_t*>(smem);                // Size: D
+    scalar_t* vr_sq_s   = reinterpret_cast<scalar_t*>(&v_s_d[D]);           // Size: R
+    scalar_t* gamma_s_d = reinterpret_cast<scalar_t*>(&vr_sq_s[R]);         // Size: D
+    
+    const scalar_t* v_b = v + bh * D;
+    scalar_t* gamma_b = gamma + bh * D;
+    
+    // Pointers for H offset
+    const scalar_t* U_h = U + h * D * R;
+    const scalar_t* W_h = W + h * D * R;
+    
+    const int tid = threadIdx.x;
+    const int bdim = blockDim.x;
+    
+    // 1. Load v into shared memory
+    for (int i = tid; i < D; i += bdim) {
+        v_s_d[i] = v_b[i];
+    }
+    __syncthreads();
+    
+    // 2. Compute v_r = v @ U -> sq = v_r^2
+    for (int r = tid; r < R; r += bdim) {
+        scalar_t sum = 0;
+        for (int j = 0; j < D; ++j) {
+            sum += v_s_d[j] * U_h[j * R + r];
+        }
+        vr_sq_s[r] = sum * sum;
+    }
+    __syncthreads();
+
+    // 3. Optional: Paper Low Rank denominator logic
+    if (is_paper_version) {
+        __shared__ scalar_t block_sum_sq;
+        if (tid == 0) block_sum_sq = 0;
+        __syncthreads();
+        
+        scalar_t local_sq = 0;
+        for (int r = tid; r < R; r += bdim) {
+            local_sq += vr_sq_s[r];
+        }
+        atomicAdd(&block_sum_sq, local_sq);
+        __syncthreads();
+        
+        scalar_t norm_vr = sqrt(block_sum_sq);
+        scalar_t denom = 1.0 + norm_vr;
+        
+        for (int r = tid; r < R; r += bdim) {
+            vr_sq_s[r] = vr_sq_s[r] / denom;
+        }
+        __syncthreads();
+    }
+    
+    // 4. Compute gamma_raw = sq @ W.T 
+    for (int d = tid; d < D; d += bdim) {
+        scalar_t sum = 0;
+        for (int r = 0; r < R; ++r) {
+            sum += vr_sq_s[r] * W_h[d * R + r];
+        }
+        gamma_s_d[d] = sum;
+    }
+    __syncthreads();
+    
+    // 5. Trace normalization (mean subtraction)
+    scalar_t mean_val = 0;
+    if (enable_trace_norm) {
+        __shared__ scalar_t block_sum_gamma;
+        if (tid == 0) block_sum_gamma = 0;
+        __syncthreads();
+        
+        scalar_t local_gamma_sum = 0;
+        for (int d = tid; d < D; d += bdim) {
+            local_gamma_sum += gamma_s_d[d];
+        }
+        atomicAdd(&block_sum_gamma, local_gamma_sum);
+        __syncthreads();
+        
+        mean_val = block_sum_gamma / D;
+    }
+    
+    // 6. Normalization and storage
+    for (int d = tid; d < D; d += bdim) {
+        scalar_t g = gamma_s_d[d];
+        if (enable_trace_norm) {
+            g -= mean_val;
+        }
+        g = clamp_val * tanh(g / clamp_val);
+        gamma_b[d] = g;
+    }
+}
+
+#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+torch::Tensor low_rank_christoffel_fwd(
+    const torch::Tensor& v, 
+    const torch::Tensor& U, 
+    const torch::Tensor& W,
+    double clamp_val,
+    bool enable_trace_norm,
+    bool is_paper_version) 
+{
+    CHECK_INPUT(v);
+    CHECK_INPUT(U);
+    CHECK_INPUT(W);
+
+    // Ensure shapes: v is [B, H, D], U is [H, D, R], W is [H, D, R]
+    int B = v.size(0);
+    int H = v.size(1);
+    int D = v.size(2);
+    int R = U.size(2);
+
+    auto gamma = torch::empty_like(v);
+
+    const int threads = 256;
+    const int blocks = B * H;
+    
+    // Shared memory size: (D + R + D) * sizeof(float)
+    const int shared_mem_size = (2 * D + R) * sizeof(float);
+
+    if (v.scalar_type() == torch::kFloat32) {
+        low_rank_christoffel_fwd_kernel<float><<<blocks, threads, shared_mem_size>>>(
+            v.data_ptr<float>(),
+            U.data_ptr<float>(),
+            W.data_ptr<float>(),
+            gamma.data_ptr<float>(),
+            B, H, D, R,
+            static_cast<float>(clamp_val),
+            enable_trace_norm,
+            is_paper_version
+        );
+    } else {
+        TORCH_CHECK(false, "low_rank_christoffel_fwd only supports float32");
+    }
+
+    return gamma;
+}

gfn/realizations/gssm/csrc/integrators/integrators.cpp

@@ -0,0 +1,252 @@
+#include <torch/extension.h>
+#include <vector>
+#define _USE_MATH_DEFINES
+#include <cmath>
+
+#ifndef M_PI
+#define M_PI 3.14159265358979323846
+#endif
+
+// -------------------------------------------------------------
+// Pure ATen Implementation of the Integrators Loop
+// -------------------------------------------------------------
+// We build this in standard C++ using ATen so it runs in a single
+// Python call but is compiled by MSVC, averting NVCC OOM bugs.
+// It performs `steps` loop using the exact GFN LowRank geometry.
+
+// Helper to compute Christoffel Gamma
+torch::Tensor _compute_gamma(
+    const torch::Tensor& v, 
+    const torch::Tensor& U, 
+    const torch::Tensor& W,
+    double clamp_val, 
+    bool enable_trace_norm, 
+    bool is_paper_version) 
+{
+    auto v_r = torch::matmul(v.unsqueeze(-2), U).squeeze(-2); // [..., R]
+    torch::Tensor sq;
+    if (is_paper_version) {
+        auto vr_norm = torch::norm(v_r, 2, -1, true);
+        sq = v_r.pow(2) / (1.0 + vr_norm);
+    } else {
+        sq = v_r.pow(2);
+    }
+    
+    auto gamma = torch::matmul(sq.unsqueeze(-2), W.transpose(-1, -2)).squeeze(-2); // [..., D]
+    
+    if (enable_trace_norm) {
+        auto mean_g = gamma.mean(-1, /*keepdim=*/true);
+        gamma = gamma - mean_g;
+    }
+    
+    return clamp_val * torch::tanh(gamma / clamp_val);
+}
+
+// Helper for velocity saturation (soft-clamp)
+torch::Tensor _clamp_velocity(const torch::Tensor& v, double v_sat) {
+    if (v_sat > 0) {
+        return v_sat * torch::tanh(v / v_sat);
+    }
+    return v;
+}
+
+// Yoshida 4th order coefficients
+const double w1 = 1.3512071919596576;
+const double w0 = -1.7024143839193153;
+const double y_c1 = w1 / 2.0;
+const double y_c2 = (w0 + w1) / 2.0;
+const double y_c3 = y_c2;
+const double y_c4 = y_c1;
+const double y_d1 = w1;
+const double y_d2 = w0;
+const double y_d3 = w1;
+
+// Helper for Gated Friction (Active Inference)
+torch::Tensor _compute_mu(
+    const torch::Tensor& x,
+    const torch::Tensor& v,
+    const torch::Tensor& gate_w,
+    const torch::Tensor& gate_b,
+    double base_friction,
+    double vel_fric_scale)
+{
+    const double eps = 1e-8;
+    const double D = x.size(-1);
+    
+    // mu_base = base_friction
+    torch::Tensor mu = torch::full_like(x.select(-1, 0).unsqueeze(-1), base_friction);
+
+    // If gate weigths are provided, calculate learnable friction component
+    if (gate_w.numel() > 0) {
+        torch::Tensor feat;
+        // Check if we need Torus features [sin, cos] (gate_w dim will be 2*D)
+        if (gate_w.size(1) == 2 * D) {
+            feat = torch::cat({torch::sin(x), torch::cos(x)}, -1); // [..., 2D]
+        } else {
+            feat = x; // Euclidean / Flat
+        }
+        
+        // Linear gate: sigmoid(feat @ w + b)
+        auto gate_out = torch::matmul(feat.unsqueeze(-2), gate_w).squeeze(-2); // [B, H, 1]
+        if (gate_b.numel() > 0) {
+            gate_out = gate_out + gate_b;
+        }
+        mu = mu + torch::sigmoid(gate_out);
+    }
+
+    // Velocity-dependent scaling: mu * (1 + scale * ||v||)
+    auto v_norm = torch::norm(v, 2, -1, true) / (std::sqrt(D) + eps);
+    mu = mu * (1.0 + vel_fric_scale * v_norm);
+    
+    return mu;
+}
+
+// Helper for Singularity Damping
+torch::Tensor _apply_singularity_damping(
+    const torch::Tensor& acc,
+    const torch::Tensor& v,
+    const torch::Tensor& U,
+    double sing_thresh,
+    double sing_strength)
+{
+    if (sing_strength <= 1.0 || sing_thresh <= 0.0) return acc;
+
+    // Detect singularity: metrics are low near singular points.
+    // In LowRank, g_diag = sum(U^2).
+    auto g_diag = (U.pow(2)).sum(-1); // [H, D]
+    
+    // Potential = sigmoid(5.0 * (g - thresh))
+    auto soft_mask = torch::sigmoid(5.0 * (g_diag - sing_thresh));
+    
+    // Scale acceleration by soft_mask (Damping Shield)
+    // Near singularity: soft_mask -> 0, damping the forces.
+    return acc * soft_mask;
+}
+
+std::vector<torch::Tensor> yoshida_fwd_aten(
+    const torch::Tensor& x_init,
+    const torch::Tensor& v_init,
+    const torch::Tensor& U,
+    const torch::Tensor& W,
+    const torch::Tensor& force,
+    const torch::Tensor& dt,
+    int steps,
+    double clamp_val,
+    double friction,
+    double vel_fric_scale,
+    double vel_sat,
+    const torch::Tensor& gate_w,
+    const torch::Tensor& gate_b,
+    double sing_thresh,
+    double sing_strength,
+    bool enable_trace_norm,
+    bool is_paper_version)
+{
+    auto x = x_init.clone();
+    auto v = v_init.clone();
+    
+    const double eps = 1e-8;
+
+    for (int i = 0; i < steps; ++i) {
+        // Sub-step 1
+        x = x + y_c1 * dt * v;
+        x = torch::remainder(x + M_PI, 2 * M_PI) - M_PI; // Toroidal resolve
+        
+        auto gamma1 = _compute_gamma(v, U, W, clamp_val, enable_trace_norm, is_paper_version);
+        auto a1_nf = force - gamma1; 
+        a1_nf = _apply_singularity_damping(a1_nf, v, U, sing_thresh, sing_strength);
+        
+        auto mu1 = _compute_mu(x, v, gate_w, gate_b, friction, vel_fric_scale);
+        
+        v = (v + y_d1 * dt * a1_nf) / (1.0 + y_d1 * dt * mu1 + eps);
+        v = _clamp_velocity(v, vel_sat);
+
+        // Sub-step 2
+        x = x + y_c2 * dt * v;
+        x = torch::remainder(x + M_PI, 2 * M_PI) - M_PI;
+        
+        auto gamma2 = _compute_gamma(v, U, W, clamp_val, enable_trace_norm, is_paper_version);
+        auto a2_nf = force - gamma2;
+        a2_nf = _apply_singularity_damping(a2_nf, v, U, sing_thresh, sing_strength);
+        
+        auto mu2 = _compute_mu(x, v, gate_w, gate_b, friction, vel_fric_scale);
+        
+        v = (v + y_d2 * dt * a2_nf) / (1.0 + y_d2 * dt * mu2 + eps);
+        v = _clamp_velocity(v, vel_sat);
+
+        // Sub-step 3
+        x = x + y_c3 * dt * v;
+        x = torch::remainder(x + M_PI, 2 * M_PI) - M_PI;
+        
+        auto gamma3 = _compute_gamma(v, U, W, clamp_val, enable_trace_norm, is_paper_version);
+        auto a3_nf = force - gamma3;
+        a3_nf = _apply_singularity_damping(a3_nf, v, U, sing_thresh, sing_strength);
+        
+        auto mu3 = _compute_mu(x, v, gate_w, gate_b, friction, vel_fric_scale);
+        
+        v = (v + y_d3 * dt * a3_nf) / (1.0 + y_d3 * dt * mu3 + eps);
+        v = _clamp_velocity(v, vel_sat);
+
+        // Final drift
+        x = x + y_c4 * dt * v;
+        x = torch::remainder(x + M_PI, 2 * M_PI) - M_PI;
+    }
+    
+    return {x, v};
+}
+
+std::vector<torch::Tensor> leapfrog_fwd_aten(
+    const torch::Tensor& x_init,
+    const torch::Tensor& v_init,
+    const torch::Tensor& U,
+    const torch::Tensor& W,
+    const torch::Tensor& force,
+    const torch::Tensor& dt,
+    int steps,
+    double clamp_val,
+    double friction,
+    double vel_fric_scale,
+    double vel_sat,
+    const torch::Tensor& gate_w,
+    const torch::Tensor& gate_b,
+    double sing_thresh,
+    double sing_strength,
+    bool enable_trace_norm,
+    bool is_paper_version)
+{
+    auto x = x_init.clone();
+    auto v = v_init.clone();
+    
+    const double eps = 1e-8;
+
+    for (int i = 0; i < steps; ++i) {
+        // Half-kick 1
+        auto gamma1 = _compute_gamma(v, U, W, clamp_val, enable_trace_norm, is_paper_version);
+        auto a1_nf = force - gamma1;
+        a1_nf = _apply_singularity_damping(a1_nf, v, U, sing_thresh, sing_strength);
+        
+        auto mu1 = _compute_mu(x, v, gate_w, gate_b, friction, vel_fric_scale);
+        
+        auto v_half = (v + 0.5 * dt * a1_nf) / (1.0 + 0.5 * dt * mu1 + eps);
+        v_half = _clamp_velocity(v_half, vel_sat);
+        
+        // Drift
+        x = x + dt * v_half;
+        x = torch::remainder(x + M_PI, 2 * M_PI) - M_PI;
+
+        // Half-kick 2
+        auto gamma2 = _compute_gamma(v_half, U, W, clamp_val, enable_trace_norm, is_paper_version);
+        auto a2_nf = force - gamma2;
+        a2_nf = _apply_singularity_damping(a2_nf, v_half, U, sing_thresh, sing_strength);
+        
+        auto mu2 = _compute_mu(x, v_half, gate_w, gate_b, friction, vel_fric_scale);
+        
+        auto a_avg = (a1_nf + a2_nf) / 2.0;
+        auto mu_avg = (mu1 + mu2) / 2.0;
+        
+        v = (v + dt * a_avg) / (1.0 + dt * mu_avg + eps);
+        v = _clamp_velocity(v, vel_sat);
+    }
+    
+    return {x, v};
+}

gfn/realizations/gssm/csrc/integrators/integrators.h

@@ -0,0 +1,41 @@
+#include <torch/extension.h>
+#include <vector>
+
+// Forward declarations for the integrators
+std::vector<torch::Tensor> yoshida_fwd_aten(
+    const torch::Tensor& x_init,
+    const torch::Tensor& v_init,
+    const torch::Tensor& U,
+    const torch::Tensor& W,
+    const torch::Tensor& force,
+    const torch::Tensor& dt,
+    int steps,
+    double clamp_val,
+    double friction,
+    double vel_fric_scale,
+    double vel_sat,
+    const torch::Tensor& gate_w,
+    const torch::Tensor& gate_b,
+    double sing_thresh,
+    double sing_strength,
+    bool enable_trace_norm,
+    bool is_paper_version);
+
+std::vector<torch::Tensor> leapfrog_fwd_aten(
+    const torch::Tensor& x_init,
+    const torch::Tensor& v_init,
+    const torch::Tensor& U,
+    const torch::Tensor& W,
+    const torch::Tensor& force,
+    const torch::Tensor& dt,
+    int steps,
+    double clamp_val,
+    double friction,
+    double vel_fric_scale,
+    double vel_sat,
+    const torch::Tensor& gate_w,
+    const torch::Tensor& gate_b,
+    double sing_thresh,
+    double sing_strength,
+    bool enable_trace_norm,
+    bool is_paper_version);

gfn/realizations/gssm/csrc/losses/toroidal.cu

@@ -0,0 +1,99 @@
+#include <torch/extension.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <math_constants.h>
+
+// ------------------------------------------------------------------------
+// Toroidal Distance Loss
+// L(y_pred, y_true) = (atan2(sin(y_pred - y_true), cos(y_pred - y_true)))^2
+// ------------------------------------------------------------------------
+
+template <typename scalar_t>
+__global__ void toroidal_distance_loss_fwd_kernel(
+    const scalar_t* __restrict__ y_pred,
+    const scalar_t* __restrict__ y_true,
+    scalar_t* __restrict__ out,
+    const int numel) 
+{
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < numel) {
+        scalar_t diff = y_pred[idx] - y_true[idx];
+        scalar_t wrapped = atan2(sin(diff), cos(diff));
+        out[idx] = wrapped * wrapped;
+    }
+}
+
+template <typename scalar_t>
+__global__ void toroidal_distance_loss_bwd_kernel(
+    const scalar_t* __restrict__ grad_output,
+    const scalar_t* __restrict__ y_pred,
+    const scalar_t* __restrict__ y_true,
+    scalar_t* __restrict__ grad_pred,
+    const int numel) 
+{
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < numel) {
+        // Derivada de atan2(sin(x), cos(x))^2 respecto a x es 2 * atan2(sin(x), cos(x))
+        scalar_t diff = y_pred[idx] - y_true[idx];
+        scalar_t wrapped = atan2(sin(diff), cos(diff));
+        grad_pred[idx] = grad_output[idx] * 2.0 * wrapped;
+    }
+}
+
+// ------------------------------------------------------------------------
+// Wrappers ATen
+// ------------------------------------------------------------------------
+
+#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+torch::Tensor toroidal_distance_loss_fwd(const torch::Tensor& y_pred, const torch::Tensor& y_true) {
+    CHECK_INPUT(y_pred);
+    CHECK_INPUT(y_true);
+
+    auto out = torch::empty_like(y_pred);
+    int numel = y_pred.numel();
+
+    const int threads = 256;
+    const int blocks = (numel + threads - 1) / threads;
+
+    if (y_pred.scalar_type() == torch::kFloat32) {
+        toroidal_distance_loss_fwd_kernel<float><<<blocks, threads>>>(
+            y_pred.data_ptr<float>(),
+            y_true.data_ptr<float>(),
+            out.data_ptr<float>(),
+            numel
+        );
+    } else {
+        TORCH_CHECK(false, "toroidal_distance_loss_fwd only supports float32");
+    }
+
+    return out;
+}
+
+torch::Tensor toroidal_distance_loss_bwd(const torch::Tensor& grad_output, const torch::Tensor& y_pred, const torch::Tensor& y_true) {
+    CHECK_INPUT(grad_output);
+    CHECK_INPUT(y_pred);
+    CHECK_INPUT(y_true);
+
+    auto grad_pred = torch::empty_like(y_pred);
+    int numel = y_pred.numel();
+
+    const int threads = 256;
+    const int blocks = (numel + threads - 1) / threads;
+
+    if (y_pred.scalar_type() == torch::kFloat32) {
+        toroidal_distance_loss_bwd_kernel<float><<<blocks, threads>>>(
+            grad_output.data_ptr<float>(),
+            y_pred.data_ptr<float>(),
+            y_true.data_ptr<float>(),
+            grad_pred.data_ptr<float>(),
+            numel
+        );
+    } else {
+        TORCH_CHECK(false, "toroidal_distance_loss_bwd only supports float32");
+    }
+
+    return grad_pred;
+}

gfn/realizations/gssm/csrc/setup.py

@@ -0,0 +1,38 @@
+import os
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+# Force sequential compilation to prevent NVCC Out-Of-Memory (LLVM ERROR)
+os.environ["MAX_JOBS"] = "1"
+
+# Directorio base
+csrc_dir = os.path.dirname(os.path.abspath(__file__))
+
+sources = [
+    os.path.join(csrc_dir, "extension.cpp"),
+    os.path.join(csrc_dir, "losses", "toroidal.cu"),
+    os.path.join(csrc_dir, "geometry", "low_rank.cu"),
+    os.path.join(csrc_dir, "integrators", "integrators.cpp")
+]
+
+# Configuración específica para MSVC/Windows vs Linux
+extra_compile_args = {
+    'cxx': ['-O2'],
+    'nvcc': ['-O2', '-allow-unsupported-compiler']
+}
+if os.name == 'nt':
+    extra_compile_args['cxx'].append('/std:c++17')
+
+setup(
+    name='gfn_cuda',
+    ext_modules=[
+        CUDAExtension(
+            name='gfn_cuda',
+            sources=sources,
+            extra_compile_args=extra_compile_args,
+        )
+    ],
+    cmdclass={
+        'build_ext': BuildExtension
+    }
+)

gfn/realizations/gssm/cuda/__init__.py

@@ -0,0 +1,11 @@
+"""
+gfn/cuda/__init__.py
+Infraestructura CUDA para GFN V5.
+"""
+import torch
+
+CUDA_AVAILABLE = torch.cuda.is_available()
+
+def is_cuda_active(tensor: torch.Tensor) -> bool:
+    """Verifica si CUDA está disponible y el tensor está en un dispositivo GPU."""
+    return CUDA_AVAILABLE and tensor.is_cuda

gfn/realizations/gssm/cuda/kernels/geometry_kernels.py

@@ -0,0 +1,99 @@
+"""
+Geometry Kernels — GFN V5
+Unified entry points for geometric computations with hardware dispatching.
+"""
+
+import torch
+from typing import Optional, Tuple, Union, Any
+from ...registry import GEOMETRY_REGISTRY
+from ...cuda import is_cuda_active
+
+# Lazy import for CUDA ops to avoid loading if not available
+_christoffel_cuda = None
+
+def _get_cuda_ops():
+    global _christoffel_cuda
+    if _christoffel_cuda is None:
+        try:
+            from ...cuda.ops import christoffel_cuda_fwd
+            _christoffel_cuda = christoffel_cuda_fwd
+        except ImportError:
+            pass
+    return _christoffel_cuda
+
+def unified_christoffel_fwd(
+    x: torch.Tensor, 
+    v: torch.Tensor, 
+    U: torch.Tensor, 
+    W: torch.Tensor, 
+    clamp_val: float = 5.0,
+    **kwargs: Any
+) -> torch.Tensor:
+    """
+    Unified forward pass for Christoffel symbols.
+    Dispatches to CUDA kernel if available and on GPU, otherwise falls back to PyTorch.
+    """
+    if is_cuda_active(v):
+        cuda_op = _get_cuda_ops()
+        if cuda_op is not None:
+            try:
+                return _run_cuda_christoffel(x, v, U, W, clamp_val, cuda_op, **kwargs)
+            except Exception as e:
+                # print(f"[Dispatcher] CUDA Error: {e}. Falling back.")
+                pass
+    
+    return _run_pytorch_christoffel(x, v, U, W, clamp_val, **kwargs)
+
+def _run_cuda_christoffel(x, v, U, W, clamp_val, cuda_op, **kwargs):
+    # Kernel expects Head-Aware tensors [B, H, HD]
+    # Check if x,v are [B, D] or [B, H, HD]
+    if v.dim() == 2:
+        x_k, v_k = x.unsqueeze(1), v.unsqueeze(1)
+    else:
+        x_k, v_k = x, v
+        
+    # U, W handling (assuming LowRank format)
+    # This logic matches the legacy kernel expectations
+    # W_k handling (ensuring rank-R is preserved per output dimension)
+    if U.dim() == 3:
+        U_k = U.transpose(1, 2).contiguous() # [H, R, HD]
+        W_k = W.transpose(1, 2).contiguous() # [H, R, HD]
+    else:
+        U_k = U.T.unsqueeze(0).contiguous() # [1, R, HD]
+        W_k = W.T.unsqueeze(0).contiguous() # [1, R, HD]
+
+    # Execute CUDA kernel
+    gamma = cuda_op(U_k, W_k, x_k, v_k, 0, 2.0, 1.0, 0.0) 
+    
+    if v.dim() == 2:
+        gamma = gamma.squeeze(1)
+    
+    return clamp_val * torch.tanh(gamma / clamp_val)
+
+def _run_pytorch_christoffel(x, v, U, W, clamp_val, **kwargs):
+    # Multi-head PyTorch fallback
+    if v.dim() == 3:
+        B, H, HD = v.shape
+        # Flatten batch and heads to use efficient matmuls
+        v_flat = v.reshape(B * H, HD)
+        if U.dim() == 3:
+            # U: [H, HD, R], W: [H, R, HD]
+            # Need to apply per-head
+            proj = torch.bmm(v.transpose(0, 1), U).transpose(0, 1) # [B, H, R]
+            sq = proj * proj
+            W_t = W.transpose(-1, -2)
+            gamma = torch.bmm(sq.transpose(0, 1), W_t).transpose(0, 1) # [B, H, HD]
+        else:
+            # Shared U, W across heads
+            proj = torch.matmul(v_flat, U) # [B*H, R]
+            sq = proj * proj
+            gamma_flat = torch.matmul(sq, W.t()) # [B*H, HD]
+            gamma = gamma_flat.view(B, H, HD)
+    else:
+        # Single head [B, D]
+        proj = torch.matmul(v, U[0] if U.dim() == 3 else U)
+        sq = proj * proj
+        W_t = (W[0] if W.dim() == 3 else W).t()
+        gamma = torch.matmul(sq, W_t)
+    
+    return clamp_val * torch.tanh(gamma / clamp_val)

gfn/realizations/gssm/cuda/kernels/integrator_kernels.py

@@ -0,0 +1,73 @@
+"""
+Integrator Kernels — GFN V5
+Unified entry points for numerical integration with hardware dispatching.
+"""
+
+import torch
+from typing import Optional, Tuple, Any, Callable
+from ...cuda import is_cuda_active
+
+# Lazy imports for CUDA kernels
+_euler_fused = None
+_rk4_fused = None
+_leapfrog_fused = None
+
+def _get_cuda_integrators():
+    global _euler_fused, _rk4_fused, _leapfrog_fused
+    if _euler_fused is None:
+        try:
+            from ...cuda.ops import euler_fused, rk4_fused, leapfrog_fused
+            _euler_fused = euler_fused
+            _rk4_fused = rk4_fused
+            _leapfrog_fused = leapfrog_fused
+        except ImportError:
+            pass
+    return _euler_fused, _rk4_fused, _leapfrog_fused
+
+def unified_leapfrog_step(
+    x: torch.Tensor, 
+    v: torch.Tensor, 
+    force: Optional[torch.Tensor],
+    U: torch.Tensor,
+    W: torch.Tensor,
+    dt: float,
+    steps: int = 1,
+    **kwargs
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Unified Leapfrog integration step."""
+    if is_cuda_active(v):
+        _, _, f_leapfrog = _get_cuda_integrators()
+        if f_leapfrog is not None:
+            try:
+                # Prep parameters for CUDA kernel
+                topo_id = kwargs.get('topology_id', 0)
+                R = kwargs.get('R', 2.0)
+                r = kwargs.get('r', 1.0)
+                H = x.shape[1] if x.dim() == 3 else 1
+                
+                # U, W transformations
+                if U.dim() == 2:
+                    # [D, R] -> [1, R, D] -> [H, R, D]
+                    U_k = U.T.unsqueeze(0).expand(H, -1, -1).contiguous()
+                else:
+                    # [H, D, R] -> [H, R, D]
+                    U_k = U.transpose(1, 2).contiguous()
+                    
+                if W.dim() == 2:
+                    # [D, R] -> [1, R] -> [H, R]
+                    # Use mean instead of sum to preserve effective force scale
+                    W_k = W.mean(dim=0).unsqueeze(0).expand(H, -1).contiguous()
+                else:
+                    # [H, D, R] -> [H, R]
+                    W_k = W.abs().mean(dim=1).contiguous()
+                
+                cx, cv = x, v
+                for _ in range(steps):
+                    cx, cv = f_leapfrog(U_k, W_k, cx, cv, force, float(dt), int(topo_id), float(R), float(r), 0.0)
+                return cx, cv
+            except Exception:
+                pass
+
+    # Python fallback is handled by the higher-level Integrator classes in gfn/integrators/
+    # This unified layer is primarily for hardware acceleration.
+    return None, None # Signal fallback

gfn/realizations/gssm/cuda/ops/__init__.py

@@ -0,0 +1,52 @@
+"""
+gfn/cuda/ops/__init__.py
+Exports all fused CUDA operations.
+Gracefully returns None for any op whose kernel is not compiled.
+"""
+from ...cuda import CUDA_AVAILABLE
+import os
+import sys
+
+# Ensure gfn/csrc is in PYTHONPATH to find the compiled .pyd
+_csrc_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "csrc"))
+if _csrc_path not in sys.path:
+    sys.path.insert(0, _csrc_path)
+
+def _get_op(module_path: str, name: str):
+    """Safely import a CUDA binding, returning None on failure."""
+    try:
+        import importlib
+        mod = importlib.import_module(module_path)
+        return getattr(mod, name, None)
+    except Exception:
+        return None
+
+# ── Geometry ──────────────────────────────────────────────────────────────────
+christoffel_cuda_fwd = _get_op("gfn_cuda", "compute_christoffel_symbols_fwd")
+christoffel_cuda_bwd = _get_op("gfn_cuda", "compute_christoffel_symbols_bwd")
+low_rank_christoffel_fwd = _get_op("gfn_cuda", "low_rank_christoffel_fwd")
+low_rank_christoffel_bwd = _get_op("gfn_cuda", "low_rank_christoffel_bwd")
+toroidal_christ_fwd = _get_op("gfn_cuda", "toroidal_geo_christoffel_fwd")
+
+# ── Integrators ───────────────────────────────────────────────────────────────
+heun_fused       = _get_op("gfn_cuda", "heun_fwd")
+leapfrog_fused   = _get_op("gfn_cuda", "leapfrog_fwd")
+yoshida_fused    = _get_op("gfn_cuda", "yoshida_fwd")
+rk4_fused        = _get_op("gfn_cuda", "rk4_fwd")
+
+# ── Loss ──────────────────────────────────────────────────────────────────────
+toroidal_loss_fwd = _get_op("gfn_cuda", "toroidal_distance_loss_fwd")
+toroidal_loss_bwd = _get_op("gfn_cuda", "toroidal_distance_loss_bwd")
+
+def __getattr__(name):
+    if name.endswith(("_fused", "_fwd", "_bwd", "_cuda")):
+        return None
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+__all__ = [
+    "CUDA_AVAILABLE",
+    "christoffel_cuda_fwd", "christoffel_cuda_bwd",
+    "low_rank_christoffel_fwd", "low_rank_christoffel_bwd",
+    "toroidal_christ_fwd", "heun_fused", "leapfrog_fused",
+    "yoshida_fused", "rk4_fused", "toroidal_loss_fwd", "toroidal_loss_bwd"
+]

gfn/realizations/gssm/data/__init__.py

@@ -0,0 +1,16 @@
+"""
+data/__init__.py — GFN V5
+"""
+from ..data.dataset import SequenceDataset
+from ..data.loader import create_dataloaders
+from ..data.transforms import shift_targets, add_bos_token, pad_sequences
+from ..data.replay import TrajectoryReplayBuffer
+
+__all__ = [
+    'SequenceDataset', 
+    'create_dataloaders', 
+    'shift_targets', 
+    'add_bos_token', 
+    'pad_sequences',
+    'TrajectoryReplayBuffer'
+]

gfn/realizations/gssm/data/dataset.py

@@ -0,0 +1,14 @@
+import torch
+from torch.utils.data import Dataset
+
+class SequenceDataset(Dataset):
+    """Simple sequence dataset for (X, Y) pairs."""
+    def __init__(self, x: torch.Tensor, y: torch.Tensor):
+        self.x = x
+        self.y = y
+        
+    def __len__(self):
+        return len(self.x)
+        
+    def __getitem__(self, idx):
+        return self.x[idx], self.y[idx]

gfn/realizations/gssm/data/loader.py

@@ -0,0 +1,53 @@
+"""
+data/loader.py — GFN V5
+DataLoaders y datasets para tareas GFN.
+WATCHOUT: El DataLoader se crea UNA vez fuera del loop de entrenamiento.
+"""
+
+import torch
+from torch.utils.data import DataLoader, Dataset, random_split
+from typing import Tuple, Optional
+from ..data.dataset import SequenceDataset
+
+
+def create_dataloaders(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    batch_size: int = 32,
+    val_split: float = 0.1,
+    shuffle: bool = True,
+    num_workers: int = 0,
+    seed: int = 42,
+) -> Tuple[DataLoader, Optional[DataLoader]]:
+    """
+    Crea train y validation DataLoaders desde tensores.
+    IMPORTANTE: Crear los DataLoaders UNA VEZ fuera del loop — no dentro.
+
+    Args:
+        x, y:        Tensores de entrada y objetivo
+        batch_size:  Tamaño de batch
+        val_split:   Fracción de datos para validación (0 = sin validación)
+        shuffle:     Mezclar datos de entrenamiento
+        num_workers: Workers para carga de datos
+        seed:        Semilla para reproducibilidad del split
+
+    Returns:
+        (train_loader, val_loader) — val_loader es None si val_split=0
+    """
+    dataset = SequenceDataset(x, y)
+
+    if val_split > 0:
+        n_val = max(1, int(len(dataset) * val_split))
+        n_train = len(dataset) - n_val
+        generator = torch.Generator().manual_seed(seed)
+        train_ds, val_ds = random_split(dataset, [n_train, n_val], generator=generator)
+
+        train_loader = DataLoader(train_ds, batch_size=batch_size,
+                                  shuffle=shuffle, num_workers=num_workers)
+        val_loader = DataLoader(val_ds, batch_size=batch_size,
+                                shuffle=False, num_workers=num_workers)
+        return train_loader, val_loader
+
+    train_loader = DataLoader(dataset, batch_size=batch_size,
+                              shuffle=shuffle, num_workers=num_workers)
+    return train_loader, None

gfn/realizations/gssm/data/replay.py

@@ -0,0 +1,130 @@
+"""
+Replay / Trajectory Buffer — GFN V5
+Maneja el almacenamiento y muestreo de estados físicos (x, v, forces) 
+para soporte de entrenamiento Off-Policy y exploración de GFlowNets reales.
+"""
+
+import torch
+from typing import Optional, Tuple
+
+class TrajectoryReplayBuffer:
+    """
+    A persistent buffer for storing and managing manifold trajectories (x, v states).
+    Serves as replay memory for Hamiltonian/Geodesic flows in V5.
+    """
+    def __init__(
+        self, 
+        capacity: int, 
+        dim: int, 
+        device: torch.device = torch.device('cpu'),
+        dtype: torch.dtype = torch.float32
+    ):
+        self.capacity = capacity
+        self.dim = dim
+        self.device = device
+        self.dtype = dtype
+        
+        # Buffers for state (x), velocity (v), and optional force
+        # Shape: [capacity, dim] or [capacity, heads, head_dim] depending on input
+        # Note: we flatten the capacity dimension but keep the geometry shape.
+        self._initialized_shape = False
+        
+        self.pointer = 0
+        self.size = 0
+        self.is_full = False
+
+    def _init_buffers(self, example_shape: torch.Size):
+        """Initializes the tensor buffers based on the first observed shape."""
+        # example_shape might be [Batch, Dim] or [Batch, Heads, HeadDim] 
+        # We need [Capacity, *shape[1:]]
+        element_shape = example_shape[1:]
+        
+        # Memory check safeguard
+        import math
+        bytes_per_el = 4 if self.dtype == torch.float32 else 8
+        total_elements = self.capacity * math.prod(element_shape)
+        # 3 buffers (x, v, force)
+        total_mb = (3 * total_elements * bytes_per_el) / (1024 ** 2)
+        if total_mb > 1024 and self.device.type == 'cuda':
+             import logging
+             logging.warning(f"ReplayBuffer: Allocating {total_mb:.1f} MB on CUDA. Risk of OOM.")
+
+        self.x_buffer = torch.zeros((self.capacity, *element_shape), device=self.device, dtype=self.dtype)
+        self.v_buffer = torch.zeros((self.capacity, *element_shape), device=self.device, dtype=self.dtype)
+        self.force_buffer = torch.zeros((self.capacity, *element_shape), device=self.device, dtype=self.dtype)
+        self._initialized_shape = True
+
+    def add(
+        self, 
+        x: torch.Tensor, 
+        v: torch.Tensor, 
+        force: Optional[torch.Tensor] = None
+    ):
+        """
+        Adds a batch of transitions to the buffer.
+        """
+        batch_size = x.size(0)
+        
+        if not self._initialized_shape:
+            self._init_buffers(x.shape)
+            
+        # Handle wrap-around indexing
+        indices = torch.arange(self.pointer, self.pointer + batch_size) % self.capacity
+        
+        self.x_buffer[indices] = x.to(self.device).detach()
+        self.v_buffer[indices] = v.to(self.device).detach()
+        if force is not None:
+            self.force_buffer[indices] = force.to(self.device).detach()
+            
+        self.pointer = (self.pointer + batch_size) % self.capacity
+        self.size = min(self.size + batch_size, self.capacity)
+        if self.size == self.capacity:
+            self.is_full = True
+
+    def sample_random(self, batch_size: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Randomly samples a batch of states from the buffer.
+        Returns: (x, v, force)
+        """
+        if self.size == 0:
+            raise ValueError("Cannot sample from an empty buffer.")
+            
+        indices = torch.randint(0, self.size, (batch_size,), device=self.device)
+        
+        return (
+            self.x_buffer[indices],
+            self.v_buffer[indices],
+            self.force_buffer[indices]
+        )
+
+    def sample_recent(self, batch_size: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Samples the most recently added transitions."""
+        if self.size == 0:
+            raise ValueError("Cannot sample from an empty buffer.")
+            
+        if self.size < batch_size:
+            idx = torch.arange(0, self.size, device=self.device)
+        else:
+            idx = (torch.arange(self.pointer - batch_size, self.pointer, device=self.device) % self.capacity)
+            
+        return (
+            self.x_buffer[idx], 
+            self.v_buffer[idx],
+            self.force_buffer[idx]
+        )
+
+    def sample_with_noise(self, batch_size: int, noise_std: float = 1e-3) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Samples with Gaussian jitter to improve robust training."""
+        x, v, _ = self.sample_random(batch_size)
+        x_noisy = x + torch.randn_like(x) * noise_std
+        return x_noisy, v
+
+    def clear(self):
+        """Resets the buffer."""
+        self.pointer = 0
+        self.size = 0
+        self.is_full = False
+        self._initialized_shape = False
+
+    def __len__(self):
+        return self.size

gfn/realizations/gssm/data/transforms.py

@@ -0,0 +1,40 @@
+"""
+data/transforms.py — GFN V5
+Transformaciones de datos para secuencias.
+"""
+
+import torch
+from typing import Tuple, Optional
+
+
+def shift_targets(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Crea pares (input, target) desplazados for language modeling.
+    input  = x[:, :-1]
+    target = x[:, 1:]
+    """
+    return x[:, :-1], x[:, 1:]
+
+
+def add_bos_token(x: torch.Tensor, bos_id: int = 0) -> torch.Tensor:
+    """Añade token BOS al inicio de cada secuencia."""
+    bos = torch.full((x.size(0), 1), bos_id, dtype=x.dtype, device=x.device)
+    return torch.cat([bos, x], dim=1)
+
+
+def pad_sequences(sequences, max_len: int, pad_id: int = 0) -> torch.Tensor:
+    """Padea una lista de secuencias de longitud variable."""
+    result = torch.full((len(sequences), max_len), pad_id, dtype=torch.long)
+    for i, seq in enumerate(sequences):
+        length = min(len(seq), max_len)
+        result[i, :length] = torch.tensor(seq[:length])
+    return result
+
+
+def create_attention_mask(lengths: torch.Tensor, max_len: int) -> torch.Tensor:
+    """
+    Crea attention mask desde longitudes de secuencia.
+    Returns: [B, max_len] con True donde hay datos válidos.
+    """
+    indices = torch.arange(max_len, device=lengths.device).unsqueeze(0)
+    return indices < lengths.unsqueeze(1)

gfn/realizations/gssm/errors.py

@@ -0,0 +1,23 @@
+class GFNError(Exception):
+    """Base exception for all GFN errors."""
+    pass
+
+class ConfigurationError(GFNError):
+    """Raised when a configuration is invalid or inconsistent."""
+    pass
+
+class GeometryError(GFNError):
+    """Raised when a geometric operation fails (e.g., out of manifold)."""
+    pass
+
+class PhysicsError(GFNError):
+    """Raised during physics engine failures (e.g., NaN detected)."""
+    pass
+
+class IntegrationError(GFNError):
+    """Raised during numerical integration failures."""
+    pass
+
+class TrainingError(GFNError):
+    """Raised during model training or optimization failures."""
+    pass

gfn/realizations/gssm/geometry/__init__.py

@@ -0,0 +1,42 @@
+"""
+gfn/geometry/__init__.py
+Public API for the geometry module — GFN V5
+"""
+
+# Base and factory
+from ..geometry.base import BaseGeometry
+from ..geometry.factory import GeometryFactory
+
+# Concrete geometries (imports trigger @register_geometry decorators)
+from ..geometry.euclidean import EuclideanGeometry
+from ..geometry.torus import ToroidalRiemannianGeometry, FlatToroidalRiemannianGeometry
+from ..geometry.low_rank import LowRankRiemannianGeometry, PaperLowRankRiemannianGeometry
+from ..geometry.adaptive import AdaptiveRiemannianGeometry
+from ..geometry.reactive import ReactiveRiemannianGeometry
+from ..geometry.hyperbolic import HyperRiemannianGeometry
+from ..geometry.holographic import HolographicRiemannianGeometry
+from ..geometry.spherical import SphericalGeometry
+from ..geometry.hierarchical import HierarchicalGeometry
+
+# Re-export FrictionGate from unified physics.components location
+from ..physics.components.friction import FrictionGate
+
+__all__ = [
+    # Base
+    "BaseGeometry",
+    "GeometryFactory",
+    # Implementations
+    "EuclideanGeometry",
+    "ToroidalRiemannianGeometry",
+    "FlatToroidalRiemannianGeometry",
+    "LowRankRiemannianGeometry",
+    "PaperLowRankRiemannianGeometry",
+    "AdaptiveRiemannianGeometry",
+    "ReactiveRiemannianGeometry",
+    "HyperRiemannianGeometry",
+    "HolographicRiemannianGeometry",
+    "SphericalGeometry",
+    "HierarchicalGeometry",
+    # Shared components
+    "FrictionGate",
+]

gfn/realizations/gssm/geometry/adaptive.py

@@ -0,0 +1,83 @@
+"""
+AdaptiveRiemannianGeometry — GFN V5
+Adaptive rank Christoffel symbol decomposition.
+Migrated from gfn/geo/riemannian/adaptive_geometry.py
+"""
+
+import torch
+import torch.nn as nn
+from typing import Optional, Union, Tuple
+
+from ..constants import CURVATURE_CLAMP
+from ..config.schema import PhysicsConfig
+from ..geometry.base import BaseGeometry
+from ..registry import register_geometry
+
+
+@register_geometry('adaptive')
+class AdaptiveRiemannianGeometry(BaseGeometry):
+    """
+    Adjusts the effective curvature rank dynamically based on velocity complexity.
+
+    Architecture:
+      eff_rank = f(||v||)  in [min_rank, max_rank]
+      Γ(v) = W[:, :eff_rank] @ (U[:, :eff_rank]^T v)^2
+    """
+
+    def __init__(self, dim: int, max_rank: int = 64, config: Optional[PhysicsConfig] = None):
+        super().__init__(config)
+        self.dim = dim
+        self.max_rank = max_rank
+        self.min_rank_ratio = 0.1
+
+        self.U_full = nn.Parameter(torch.randn(dim, max_rank) * 0.01)
+        self.W_full = nn.Parameter(torch.randn(dim, max_rank) * 0.01)
+
+        # Complexity predictor: maps v → rank_ratio ∈ [0, 1]
+        self.complexity_net = nn.Sequential(
+            nn.Linear(dim, 32),
+            nn.ReLU(),
+            nn.Linear(32, 1),
+            nn.Sigmoid()
+        )
+        # Initialize bias to start with a mostly-open rank to avoid vanishing gradients
+        nn.init.constant_(self.complexity_net[-2].bias, 1.0)
+        
+        self.return_friction_separately = True
+
+    def forward(self, x: torch.Tensor, v: Optional[torch.Tensor] = None,
+                force: Optional[torch.Tensor] = None, **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        if v is None:
+            return torch.zeros_like(x)
+
+        # 1. Predict rank weight (soft-mask)
+        # Use complexity_net to predict a value p in [0, 1]
+        p = self.complexity_net(v)  # [B, 1]
+        
+        # Create a soft mask for the rank dimension: [B, max_rank]
+        # Mask[i] = sigmoid(slope * (p * max_rank - i))
+        # This approximates hard-slicing but is differentiable.
+        indices = torch.arange(self.max_rank, device=v.device).float()
+        slope = 10.0
+        soft_mask = torch.sigmoid(slope * (p * self.max_rank - indices)) # [B, max_rank]
+
+        # 2. Christoffel using all components modulated by mask
+        proj = torch.matmul(v, self.U_full)   # [B, max_rank]
+        sq = proj * proj                     # [B, max_rank]
+        modulated_sq = sq * soft_mask         # [B, max_rank]
+        gamma = torch.matmul(modulated_sq, self.W_full.t())  # [B, dim]
+
+        # 3. Friction (ensure mu is not just zero)
+        # Fallback to config friction or a base value
+        friction_base = getattr(self.config.stability, 'friction', 0.1)
+        mu = torch.full_like(v, friction_base)
+
+        gamma_clamped = CURVATURE_CLAMP * torch.tanh(gamma / CURVATURE_CLAMP)
+
+        if self.return_friction_separately:
+            return gamma_clamped, mu
+
+        return gamma_clamped + mu * v
+
+    def metric_tensor(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.ones_like(x)

gfn/realizations/gssm/geometry/base.py

@@ -0,0 +1,70 @@
+import torch
+import torch.nn as nn
+from typing import Optional, Tuple, Union
+from ..interfaces.geometry import Geometry
+from ..config.schema import PhysicsConfig
+from ..constants import TOPOLOGY_EUCLIDEAN
+
+class BaseGeometry(nn.Module):
+    """
+    Base implementation for Riemannian Geometries in GFN V5.
+    Conforms to the Geometry protocol.
+    """
+    def __init__(self, config: Optional[PhysicsConfig] = None):
+        super().__init__()
+        self.config = config or PhysicsConfig()
+        self.return_friction_separately = True
+        self.topology_type = self.config.topology.type
+        
+    def metric_tensor(self, x: torch.Tensor) -> torch.Tensor:
+        """Default metric is identity (Euclidean). Subclasses should override."""
+        return torch.ones_like(x)
+        
+    def christoffel_symbols(self, x: torch.Tensor) -> torch.Tensor:
+        """Default Christoffel symbols are zero (Euclidean). Subclasses should override."""
+        return torch.zeros_like(x)
+
+    def compute_kinetic_energy(self, x: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+        """
+        Calculates Riemannian kinetic energy: T = (1/2) Σ_i g_ii v_i²
+        Supports position-dependent metrics (like Torus).
+        """
+        g = self.metric_tensor(x)  # [..., D]
+        return 0.5 * (g * v.pow(2)).sum(dim=-1)
+
+    def compute_potential_energy(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Calculates physical potential energy V(x). 
+        Default is 0.0 unless overwritten by specific topologies or forces.
+        """
+        return torch.zeros_like(x).sum(dim=-1)
+        
+    def forward(self, x: torch.Tensor, v: Optional[torch.Tensor] = None, force: Optional[torch.Tensor] = None) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        """
+        Computes acceleration: acc = -Gamma(v, v) + F/g
+        Subclasses can override for more complex physics.
+        """
+        if v is None:
+            return torch.zeros_like(x)
+            
+        gamma = self.christoffel_symbols(x)
+        # Standard geodesic acceleration: -Gamma^k_ij v^i v^j
+        # In our simplified 1D-per-dimension metric, it's often just a point-wise product
+        acc = -gamma * (v**2) 
+        
+        if force is not None:
+            g = self.metric_tensor(x)
+            acc = acc + (force / (g + 1e-8))
+            
+        if getattr(self, 'return_friction_separately', False):
+            return acc, torch.zeros_like(v) if v is not None else torch.zeros_like(x)
+            
+        return acc
+
+    def project(self, x: torch.Tensor) -> torch.Tensor:
+        """Default projection is identity. Subclasses should override for periodic spaces."""
+        return x
+
+    def dist(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
+        """Default distance is Euclidean. Subclasses should override."""
+        return torch.norm(x1 - x2, dim=-1)

gfn/realizations/gssm/geometry/euclidean.py

@@ -0,0 +1,20 @@
+import torch
+from .base import BaseGeometry
+from ..registry import register_geometry
+from ..constants import TOPOLOGY_EUCLIDEAN
+
+@register_geometry(TOPOLOGY_EUCLIDEAN)
+class EuclideanGeometry(BaseGeometry):
+    """Standard Euclidean Space (Flat)."""
+    
+    def metric_tensor(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.ones_like(x)
+        
+    def christoffel_symbols(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.zeros_like(x)
+        
+    def project(self, x: torch.Tensor) -> torch.Tensor:
+        return x
+
+    def dist(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
+        return torch.norm(x1 - x2, dim=-1)

gfn/realizations/gssm/geometry/factory.py

@@ -0,0 +1,117 @@
+"""
+GeometryFactory — GFN V5
+Creates geometry instances from PhysicsConfig.
+Supports: euclidean, torus, low_rank, reactive, adaptive, hyperbolic, holographic.
+"""
+
+from typing import Optional
+from ..config.schema import PhysicsConfig
+from ..registry import GEOMETRY_REGISTRY
+from ..constants import TOPOLOGY_TORUS, TOPOLOGY_EUCLIDEAN
+import logging
+
+logger = logging.getLogger(__name__)
+
+_GEOMETRIES_REGISTERED = False
+
+def _register_all_geometries():
+    """Importa los submódulos explícitamente para registrar las geometrías."""
+    global _GEOMETRIES_REGISTERED
+    if _GEOMETRIES_REGISTERED:
+        return
+    from . import euclidean
+    from . import torus
+    from . import low_rank
+    from . import adaptive
+    from . import reactive
+    from . import hyperbolic
+    _GEOMETRIES_REGISTERED = True
+
+class GeometryFactory:
+    """
+    Creates manifold geometries from configuration.
+
+    Primary key: topology.type  ('euclidean', 'torus', 'hyperbolic', ...)
+    Secondary key: topology.riemannian_type  ('low_rank', 'reactive', 'adaptive', ...)
+
+    riemannian_type overrides topology.type when explicitly set and registered.
+    """
+
+    @staticmethod
+    def _lookup_key(config: PhysicsConfig) -> str:
+        _register_all_geometries()
+        topo_type = config.topology.type.lower()
+        riem_type = getattr(config.topology, 'riemannian_type', 'reactive').lower()
+        available = GEOMETRY_REGISTRY.list_keys()
+        
+        # Priority Logic:
+        # 1. Prioritize learned Riemannian geometries (low_rank, reactive, adaptive)
+        #    even if the topology is specialized (torus, etc.), as they handle topology via features.
+        learned_types = {'low_rank', 'reactive', 'adaptive', 'low_rank_paper'}
+        if riem_type in learned_types and riem_type in available:
+            return riem_type
+
+        # 2. Otherwise, if topology is specific (torus, hyperbolic, etc.), use its analytical model.
+        if topo_type in available and topo_type != TOPOLOGY_EUCLIDEAN:
+            return topo_type
+            
+        # 3. Fallback to riem_type or topo_type
+        if riem_type in available:
+            return riem_type
+            
+        return topo_type
+
+    @staticmethod
+    def create(config: PhysicsConfig):
+        """
+        Create geometry using default dim from config.
+        Looks for 'dim' in topology config or falls back to 64.
+        """
+        lookup_key = GeometryFactory._lookup_key(config)
+        available = GEOMETRY_REGISTRY.list_keys()
+
+        if lookup_key in available:
+            geometry_cls = GEOMETRY_REGISTRY.get(lookup_key)
+            try:
+                dim = getattr(config, 'dim', 64)
+                rank = getattr(config.topology, 'riemannian_rank', 16)
+                return geometry_cls(dim=dim, rank=rank, config=config)
+            except TypeError:
+                try:
+                    return geometry_cls(config=config)
+                except TypeError:
+                    return geometry_cls()
+
+        logger.warning(f"Geometry '{lookup_key}' not found. Using EuclideanGeometry.")
+        from .euclidean import EuclideanGeometry
+        return EuclideanGeometry(config=config)
+
+    @staticmethod
+    def create_with_dim(dim: int, rank: int, num_heads: int, config: PhysicsConfig):
+        """
+        Create geometry with explicit dim and rank.
+        Used by ModelFactory to pass head_dim (not total dim) to the geometry,
+        since geometry operates on per-head tensors [B, H, HD].
+        """
+        lookup_key = GeometryFactory._lookup_key(config)
+        available = GEOMETRY_REGISTRY.list_keys()
+
+        if lookup_key in available:
+            geometry_cls = GEOMETRY_REGISTRY.get(lookup_key)
+            try:
+                return geometry_cls(dim=dim, rank=rank, num_heads=num_heads, config=config)
+            except TypeError:
+                try:
+                    return geometry_cls(dim=dim, rank=rank, config=config)
+                except TypeError:
+                    try:
+                         return geometry_cls(config=config)
+                    except TypeError:
+                         return geometry_cls()
+
+        logger.warning(f"Geometry '{lookup_key}' not found. Using EuclideanGeometry.")
+        from .euclidean import EuclideanGeometry
+        try:
+             return EuclideanGeometry(dim=dim, num_heads=num_heads, config=config)
+        except TypeError:
+             return EuclideanGeometry(config=config)

gfn/realizations/gssm/geometry/hierarchical.py

@@ -0,0 +1,84 @@
+import torch
+import torch.nn as nn
+from typing import List, Optional, Union, Tuple, Any
+from ..geometry.base import BaseGeometry
+from ..geometry.low_rank import LowRankRiemannianGeometry
+from ..registry import register_geometry
+
+@register_geometry('hierarchical')
+class HierarchicalGeometry(BaseGeometry):
+    """
+    Multi-Scale Riemannian Geometry (Christoffel Mixture).
+    Combines multiple geometries (typically Low-Rank) with different scales.
+    
+    Migrated from gfn_old HierarchicalRiemannianGeometry.
+    """
+    def __init__(self, dim: int, rank: int = 16, ranks: Optional[List[int]] = None, 
+                 num_heads: int = 1, config: Optional[Any] = None, **kwargs):
+        super().__init__(config)
+        self.dim = dim
+        self.ranks = ranks if ranks is not None else [8, 16, 32]
+        if rank not in self.ranks:
+            # Optionally include the factory-suggested rank
+            self.ranks = sorted(list(set(self.ranks + [rank])))
+        self.num_heads = num_heads
+        
+        # Initialize sub-geometries (defaulting to LowRank)
+        self.scales = nn.ModuleList([
+            LowRankRiemannianGeometry(dim, rank=r, num_heads=num_heads, config=config)
+            for r in self.ranks
+        ])
+        
+        # Learnable mixing weights
+        self.scale_weights = nn.Parameter(torch.ones(len(self.ranks)) / len(self.ranks))
+        self.return_friction_separately = False
+
+    def forward(self, x: torch.Tensor, v: Optional[torch.Tensor] = None, 
+                force: Optional[torch.Tensor] = None, **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        
+        gammas = []
+        frictions = []
+        
+        # Execute each scale
+        for scale in self.scales:
+            # Temporarily ensure consistent return mode
+            was_sep = getattr(scale, 'return_friction_separately', False)
+            scale.return_friction_separately = True
+            
+            res = scale(x, v, force=force, **kwargs)
+            if isinstance(res, tuple):
+                g, f = res
+            else:
+                g, f = res, torch.zeros_like(v) if v is not None else torch.zeros_like(x)
+            
+            gammas.append(g)
+            frictions.append(f)
+            scale.return_friction_separately = was_sep
+
+        # Mix using softmax weights
+        weights = torch.softmax(self.scale_weights, dim=0)
+        
+        gamma_mixed = sum(w * g for w, g in zip(weights, gammas))
+        friction_mixed = sum(w * f for w, f in zip(weights, frictions))
+        
+        if self.return_friction_separately:
+            return gamma_mixed, friction_mixed
+            
+        if v is not None:
+             return gamma_mixed + friction_mixed * v
+        return gamma_mixed
+
+    def metric_tensor(self, x: torch.Tensor) -> torch.Tensor:
+        weights = torch.softmax(self.scale_weights, dim=0)
+        metrics = [scale.metric_tensor(x) for scale in self.scales]
+        return sum(w * m for w, m in zip(weights, metrics))
+
+    def project(self, x: torch.Tensor) -> torch.Tensor:
+        # Hierarchical usually just projects using the first scale (geometry consistent)
+        from typing import cast
+        return cast(BaseGeometry, self.scales[0]).project(x)
+
+    def dist(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
+        weights = torch.softmax(self.scale_weights, dim=0)
+        dists = [scale.dist(x1, x2) for scale in self.scales]
+        return sum(w * d for w, d in zip(weights, dists))

gfn/realizations/gssm/geometry/holographic.py

@@ -0,0 +1,91 @@
+"""
+HolographicRiemannianGeometry — GFN V5
+AdS/CFT-inspired holographic extensions (Paper 18).
+Migrated from gfn/geo/physical/holographic_geometry.py
+"""
+
+import torch
+import torch.nn as nn
+from typing import Optional, Union, Tuple
+
+from ..config.schema import PhysicsConfig
+from ..geometry.base import BaseGeometry
+from ..registry import register_geometry
+
+
+@register_geometry('holographic')
+class HolographicRiemannianGeometry(BaseGeometry):
+    """
+    Conformal manifold inspired by Bulk-Boundary (AdS/CFT) correspondence.
+
+    Lifts boundary state x → bulk (x, z) where z is the holographic radial dim.
+    Conformal metric: g_ij = (1/z(x)²) · δ_ij
+
+    The Christoffel correction adds an AdS-geodesic term to any base geometry.
+    """
+
+    def __init__(self, base_geometry: BaseGeometry, z_min: float = 0.1,
+                 z_max: float = 10.0, config: Optional[PhysicsConfig] = None):
+        super().__init__(config)
+        self.base_geometry = base_geometry
+        self.dim = getattr(base_geometry, 'dim', None)
+        self.z_min = z_min
+        self.z_max = z_max
+
+        dim = self.dim or 0
+        if dim > 0:
+            self.radial_net: nn.Module = nn.Sequential(
+                nn.Linear(dim, dim // 2),
+                nn.SiLU(),
+                nn.Linear(dim // 2, 1),
+                nn.Softplus()
+            )
+        else:
+            self.radial_net = nn.Identity()
+
+    def get_z_and_grad(self, x: torch.Tensor):
+        x_req = x.detach().requires_grad_(True)
+        with torch.enable_grad():
+            z = self.radial_net(x_req) + self.z_min
+            z = torch.clamp(z, max=self.z_max)
+            grad_z = torch.autograd.grad(
+                z.sum(), x_req,
+                create_graph=self.training,
+                retain_graph=False
+            )[0]
+        return z, grad_z
+
+    def metric_tensor(self, x: torch.Tensor) -> torch.Tensor:
+        z, _ = self.get_z_and_grad(x)
+        return (1.0 / z.pow(2)) * torch.ones_like(x)
+
+    def forward(self, x: torch.Tensor, v: Optional[torch.Tensor] = None,
+                force: Optional[torch.Tensor] = None, **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        out_base = self.base_geometry(x, v, force=force, **kwargs)
+        if isinstance(out_base, tuple):
+            gamma_base, mu = out_base
+        else:
+            gamma_base, mu = out_base, torch.zeros_like(v) if v is not None else torch.zeros_like(x)
+
+        if v is None:
+            if self.return_friction_separately:
+                return gamma_base, mu
+            return gamma_base
+
+        z, grad_z = self.get_z_and_grad(x)
+        v_dot_gradz = (v * grad_z).sum(dim=-1, keepdim=True)
+        v_sq = (v * v).sum(dim=-1, keepdim=True)
+        gamma_ads = -(1.0 / z) * (2.0 * v_dot_gradz * v - v_sq * grad_z)
+
+        gamma_total = gamma_base + gamma_ads
+        
+        if self.return_friction_separately:
+            return gamma_total, mu
+            
+        return gamma_total + mu * v
+
+    def project(self, x: torch.Tensor) -> torch.Tensor:
+        return self.base_geometry.project(x)
+
+    def dist(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
+        return self.base_geometry.dist(x1, x2)

gfn/realizations/gssm/geometry/hyperbolic.py

@@ -0,0 +1,97 @@
+"""
+HyperRiemannianGeometry — GFN V5
+Context-dependent (gated) Christoffel symbols.
+Migrated from gfn/geo/topological/hyperbolic_geometry.py
+"""
+
+import torch
+import torch.nn as nn
+from typing import Optional, Union, Tuple
+
+from ..constants import CURVATURE_CLAMP, EPS, TOPOLOGY_TORUS
+from ..config.schema import PhysicsConfig
+from ..geometry.low_rank import LowRankRiemannianGeometry
+from ..registry import register_geometry
+
+
+@register_geometry('hyperbolic')
+class HyperRiemannianGeometry(LowRankRiemannianGeometry):
+    """
+    Hyper-Christoffel: geometry conditioned on current position.
+
+    Architecture:
+      U(x) = U_static * diag(Gate_u(x))   — position-scaled basis
+      W(x) = W_static * diag(Gate_w(x))
+      Γ(v | x) = W(x) @ (U(x)^T v)²
+
+    Gates output values in [0, 2] initialized near 1.0 (identity).
+    """
+
+    def __init__(self, dim: int, rank: int = 16, num_heads: int = 1,
+                 config: Optional[PhysicsConfig] = None):
+        super().__init__(dim, rank, num_heads=num_heads, config=config)
+        self.return_friction_separately = True
+
+        self.gate_u = nn.Linear(dim, rank)
+        self.gate_w = nn.Linear(dim, rank)
+        nn.init.zeros_(self.gate_u.weight); nn.init.zeros_(self.gate_u.bias)
+        nn.init.zeros_(self.gate_w.weight); nn.init.zeros_(self.gate_w.bias)
+
+    def forward(self, x: torch.Tensor, v: Optional[torch.Tensor] = None,
+                force: Optional[torch.Tensor] = None, **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        if v is None:
+            return torch.zeros_like(x)
+
+        original_shape = v.shape
+        # Handle multi-head [B, H, HD] -> [B*H, HD]
+        if v.dim() == 3:
+            B, H, HD = v.shape
+            v_flat = v.reshape(B * H, HD)
+            x_flat = x.reshape(B * H, HD)
+        else:
+            v_flat = v
+            x_flat = x
+            B, H = None, None
+
+        # Context gates in [0, 2]
+        g_u = torch.sigmoid(self.gate_u(x_flat)) * 2.0  # [B*H, rank]
+        g_w = torch.sigmoid(self.gate_w(x_flat)) * 2.0
+
+        # Modulate static basis
+        # self.U is [HD, rank] or [H, HD, rank]
+        U_eff = self.U if self.U.dim() == 2 else self.U.mean(0)
+        proj_static = torch.matmul(v_flat, U_eff) # [B*H, rank]
+        proj_dynamic = proj_static * g_u
+
+        # Soft-saturation to prevent energy explosion
+        sq_dynamic = (proj_dynamic * proj_dynamic) / (1.0 + torch.abs(proj_dynamic) + EPS)
+        sq_modulated = sq_dynamic * g_w
+
+        W_t = self.W.t() if self.W.dim() == 2 else self.W.mean(0).t()
+        gamma = torch.matmul(sq_modulated, W_t) # [B*H, HD]
+
+        # Restore original shape if multi-head
+        if B is not None:
+            gamma = gamma.view(original_shape)
+            x_flat_for_mu = x_flat
+            v_flat_for_mu = v_flat
+        else:
+            x_flat_for_mu = x
+            v_flat_for_mu = v
+
+        # Friction
+        x_in = torch.cat([torch.sin(x_flat), torch.cos(x_flat)], dim=-1) \
+            if self.topology_type == TOPOLOGY_TORUS else x_flat
+        mu_base = self.friction + self.friction_gate(x_in, force=force)
+        v_norm = torch.norm(v, dim=-1, keepdim=True) / (self.dim ** 0.5 + EPS)
+        mu = mu_base * (1.0 + self.velocity_friction_scale * v_norm)
+        if mu.shape != v.shape:
+            mu = mu.view_as(v) if mu.numel() == v.numel() else mu.mean(dim=-1, keepdim=True)
+
+        gamma = self._normalize(gamma)
+        gamma = self.clamp_val * torch.tanh(gamma / self.clamp_val)
+
+        if self.return_friction_separately:
+            return gamma, mu
+
+        return gamma + mu * v

gfn/realizations/gssm/geometry/low_rank.py

@@ -0,0 +1,324 @@
+"""
+LowRankRiemannianGeometry — GFN V5
+Computes Christoffel symbols via a low-rank decomposition.
+Migrated from gfn/geo/riemannian/low_rank_geometry.py
+"""
+
+import torch
+import torch.nn as nn
+from typing import Optional, Union, Tuple, Dict, Any
+
+from ..constants import (
+    EPS, MAX_VELOCITY, TOPOLOGY_TORUS, TOPOLOGY_EUCLIDEAN,
+    DEFAULT_FRICTION, CURVATURE_CLAMP, GATE_BIAS_OPEN
+)
+from ..config.schema import PhysicsConfig
+from ..geometry.base import BaseGeometry
+from ..registry import register_geometry
+from ..cuda.ops import CUDA_AVAILABLE, low_rank_christoffel_fwd, low_rank_christoffel_bwd
+
+class LowRankChristoffelFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, v, U, W, clamp_val, enable_trace_norm, is_paper_version=False):
+        v_c = v.contiguous()
+        U_c = U.contiguous()
+        W_c = W.contiguous()
+        
+        gamma = low_rank_christoffel_fwd(v_c, U_c, W_c, float(clamp_val), enable_trace_norm, is_paper_version)
+        ctx.save_for_backward(v_c, U_c, W_c, gamma)
+        ctx.clamp_val = float(clamp_val)
+        ctx.enable_trace_norm = enable_trace_norm
+        ctx.is_paper_version = is_paper_version
+        return gamma
+
+    @staticmethod
+    def backward(ctx, grad_gamma):
+        v_c, U_c, W_c, gamma_out = ctx.saved_tensors
+        if grad_gamma is None:
+            return None, None, None, None, None, None
+            
+        grad_gamma_c = grad_gamma.contiguous()
+        d_v, d_U, d_W = low_rank_christoffel_bwd(
+            grad_gamma_c, v_c, U_c, W_c, gamma_out, 
+            ctx.clamp_val, ctx.enable_trace_norm, ctx.is_paper_version
+        )
+        return d_v, d_U, d_W, None, None, None
+
+
+# Use unified FrictionGate from physics.components (no duplication)
+from ..physics.components.friction import FrictionGate
+
+
+@register_geometry('low_rank')
+class LowRankRiemannianGeometry(BaseGeometry):
+    r"""
+    Low-rank Christoffel symbol decomposition.
+
+    Γ^k_ij ≈ Σ_r W_{rk} * (U_ir * U_jr)
+    This is an approximation — symmetry is preserved but Bianchi identities are not guaranteed.
+    Chosen for computational efficiency.
+
+    Args:
+        dim: Manifold dimension.
+        rank: Rank of the decomposition.
+        num_heads: Number of parallel heads.
+        config: PhysicsConfig instance.
+    """
+
+    def __init__(self, dim: int, rank: int = 16, num_heads: int = 1,
+                 config: Optional[PhysicsConfig] = None):
+        super().__init__(config)
+        self.dim = dim
+        self.rank = rank
+        self.num_heads = num_heads
+
+        topo = self.config.topology.type.lower()
+        self.topology = topo
+        self.clamp_val = self.config.stability.curvature_clamp
+        self.enable_trace_normalization = self.config.stability.enable_trace_normalization
+        self.enable_trace_normalization = self.config.stability.enable_trace_normalization
+        # Friction parameters are now handled by PhysicsEngine to avoid duplication
+
+        # Feature dimension for gate input (Fourier for torus)
+        gate_input_dim = dim * 2 if topo == TOPOLOGY_TORUS else dim
+
+        # Low-rank basis parameters - initialized with small noise to break symmetry
+        if num_heads > 1:
+            self.U = nn.Parameter(torch.randn(num_heads, dim, rank) * 1e-4)
+            self.W = nn.Parameter(torch.randn(num_heads, dim, rank) * 1e-4)
+        else:
+            self.U = nn.Parameter(torch.randn(dim, rank) * 1e-4)
+            self.W = nn.Parameter(torch.randn(dim, rank) * 1e-4)
+
+        # Friction gate
+        friction_mode = getattr(self.config.stability, 'friction_mode', 'static')
+        self.friction_gate = FrictionGate(dim, gate_input_dim, mode=friction_mode, num_heads=num_heads)
+
+        # CONTRACT: LowRank ALWAYS returns (gamma_christoffel, mu_friction) separately.
+        # The physics engine is the single authority on when/how friction is applied.
+        # This prevents the P0.1 double-friction bug (geometry + engine both applying mu*v).
+        self.return_friction_separately = True
+
+    def _get_features(self, x: torch.Tensor) -> torch.Tensor:
+        """Convert position to gate input features."""
+        if self.topology == TOPOLOGY_TORUS:
+            return torch.cat([torch.sin(x), torch.cos(x)], dim=-1)
+        return x
+
+    def connection(self, v: torch.Tensor, w: torch.Tensor,
+                   x: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Bilinear Christoffel contraction: Γ(v,w)^k
+        Γ^k_ij ≈ Σ_r W[r,k] * (U[i,r] * U[j,r])
+        """
+        if v.dim() == 3 and self.U.dim() == 3:
+            v_r = torch.einsum('bhd, hdr -> bhr', v, self.U)
+            w_r = torch.einsum('bhd, hdr -> bhr', w, self.U)
+            vw_r = v_r * w_r
+            gamma = torch.einsum('bhr, hdr -> bhd', vw_r, self.W)
+        else:
+            v_r = v @ self.U   # [..., rank]  (works for both 2D and 3D U)
+            w_r = w @ self.U
+            vw_r = v_r * w_r
+            W_t = self.W.transpose(-1, -2) if self.W.dim() == 3 else self.W.t()
+            gamma = vw_r @ W_t
+        return torch.clamp(gamma, -self.clamp_val, self.clamp_val)
+
+    def _normalize(self, gamma: torch.Tensor) -> torch.Tensor:
+        """Symmetry-preserving trace normalization."""
+        if gamma.dim() < 2:
+            return gamma
+
+        is_multi_head = (gamma.dim() == 3 and self.num_heads > 1)
+
+        # Matrix case [..., D, D]
+        if not is_multi_head and gamma.dim() >= 3 and gamma.shape[-1] == gamma.shape[-2]:
+            gamma_sym = 0.5 * (gamma + gamma.transpose(-1, -2))
+            if self.enable_trace_normalization:
+                diag_mean = torch.diagonal(gamma_sym, dim1=-1, dim2=-2).mean(dim=-1, keepdim=True)
+                correction = torch.diag_embed(diag_mean.expand(-1, self.dim))
+                return gamma_sym - correction
+            return gamma_sym
+
+        # Vector case [..., D]
+        if self.enable_trace_normalization:
+            return gamma - gamma.mean(dim=-1, keepdim=True)
+        return gamma
+
+    def forward(self, x: torch.Tensor, v: Optional[torch.Tensor] = None,
+                force: Optional[torch.Tensor] = None, **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        if v is None:
+            return torch.zeros_like(x)
+
+        original_shape = v.shape
+        # Handle multi-head [B, H, HD] → reshape to [B*H, HD] for matmul with U=[HD, rank]
+        if v.dim() == 3 and self.U.dim() == 2:
+            B, H, HD = v.shape
+            v_flat = v.reshape(B * H, HD)   # [B*H, HD]
+            x_flat = x.reshape(B * H, HD)
+        else:
+            v_flat = v
+            x_flat = x
+        B, H, HD = None, None, v_flat.shape[-1]
+        R = self.rank
+        
+        # Check if we can take the fast CUDA path
+        use_cuda_fused = (
+            CUDA_AVAILABLE and 
+            low_rank_christoffel_fwd is not None and 
+            v_flat.is_cuda and 
+            v_flat.dtype == torch.float32 and
+            self.W.dim() == 3
+        )
+
+        if use_cuda_fused:
+            # Reshape [B*H, HD] -> [B, H, HD] to match what kernel expects
+            actual_B = original_shape[0] if v.dim() == 3 else 1
+            actual_H = self.num_heads
+            
+            v_re = v_flat.view(actual_B, actual_H, HD)
+            U_re = self.U.view(actual_H, HD, R)  # self.U is [H, D, R]
+            W_re = self.W.view(actual_H, HD, R)
+            
+            gamma_re = LowRankChristoffelFunction.apply(
+                v_re, U_re, W_re, self.clamp_val, self.enable_trace_normalization, False
+            )
+            gamma = gamma_re.view_as(v_flat)
+        else:
+            # Christoffel symbols via self-connection (Native Fallback)
+            if v_flat.dim() == 3 and self.U.dim() == 3:
+                v_r = torch.einsum('bhd, hdr -> bhr', v_flat, self.U)
+                sq = v_r * v_r
+                gamma = torch.einsum('bhr, hdr -> bhd', sq, self.W)
+            else:
+                v_r = v_flat @ self.U    # [..., rank]
+                sq = v_r * v_r
+                W_t = self.W.transpose(-1, -2) if self.W.dim() == 3 else self.W.t()
+                gamma = sq @ W_t         # [..., HD]
+
+        # Friction coefficient (position-dependent, gated)
+        x_in = self._get_features(x_flat)
+        mu = self.friction_gate(x_in, force=force)
+        
+        # Note: PhysicsEngine will add base friction and apply velocity scaling
+
+        # Friction and normalizing only if not already done in CUDA
+        if not use_cuda_fused:
+            gamma = self._normalize(gamma)
+            gamma = self.clamp_val * torch.tanh(gamma / self.clamp_val)
+
+        # Restore original shape if we reshaped
+        if B is not None:
+            gamma = gamma.view(original_shape)
+            mu = mu.view(original_shape)
+
+        # CONTRACT: always return (gamma_pure, mu) so engine has single authority over friction
+        return gamma, mu
+
+    def metric_tensor(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Implicit Riemannian metric from the low-rank decomposition.
+        
+        The Christoffel parametrization Γ^k_ij ≈ Σ_r W_rk (U_ir U_jr) implies
+        an underlying metric: g_ij ≈ Σ_r U_ir * U_jr = diag(U @ Uᵀ)
+        
+        Returns per-coordinate metric scale [..., D] or ones if shape unknown.
+        T = (1/2) Σ_i g_ii v_i²  (Riemannian kinetic energy)
+        """
+        if self.U.dim() == 2:
+            # Single head: U is [D, rank] → g_diag is [D]
+            g_diag = (self.U ** 2).sum(dim=-1)  # [D]
+            # Broadcast to x shape: handles [B, D], [B*H, D], any [..., D]
+            return g_diag.expand_as(x)
+        else:
+            # Multi-head: U is [H, D, rank] → g_diag is [H, D]
+            g_diag = (self.U ** 2).sum(dim=-1)  # [H, D]
+            if x.dim() == 3 and x.shape[1] == self.num_heads:
+                # [B, H, D]: structured multi-head
+                return g_diag.unsqueeze(0).expand_as(x)
+            else:
+                # [B, H*D]: flat layout — expand g_diag [H,D] → [H*D], broadcast to [B, H*D]
+                g_flat = g_diag.reshape(-1)  # [H*D]
+                return g_flat.expand(x.shape[0], -1) if x.dim() == 2 else g_flat.expand_as(x)
+
+    def dist(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
+        if self.topology == TOPOLOGY_TORUS:
+            diff = x1 - x2
+            diff = torch.atan2(torch.sin(diff), torch.cos(diff))
+            return torch.norm(diff, dim=-1)
+        return torch.norm(x1 - x2, dim=-1)
+
+    def project(self, x: torch.Tensor) -> torch.Tensor:
+        if self.topology == TOPOLOGY_TORUS:
+            return torch.atan2(torch.sin(x), torch.cos(x))
+        return x
+
+
+@register_geometry('low_rank_paper')
+class PaperLowRankRiemannianGeometry(LowRankRiemannianGeometry):
+    def __init__(self, dim: int, rank: int = 16, num_heads: int = 1,
+                 config: Optional[PhysicsConfig] = None):
+        super().__init__(dim, rank, num_heads=num_heads, config=config)
+        self.return_friction_separately = True
+
+    def forward(self, x: torch.Tensor, v: Optional[torch.Tensor] = None,
+                force: Optional[torch.Tensor] = None, **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        if v is None:
+            return torch.zeros_like(x)
+
+        original_shape = v.shape
+        if v.dim() == 3 and self.U.dim() == 2:
+            B, H, HD = v.shape
+            v_flat = v.reshape(B * H, HD)
+            x_flat = x.reshape(B * H, HD)
+        else:
+            v_flat = v
+            x_flat = x
+        B, H, HD = None, None, v_flat.shape[-1]
+        R = self.rank
+
+        use_cuda_fused = (
+            CUDA_AVAILABLE and 
+            low_rank_christoffel_fwd is not None and 
+            v_flat.is_cuda and 
+            v_flat.dtype == torch.float32 and
+            self.W.dim() == 3
+        )
+
+        if use_cuda_fused:
+            actual_B = original_shape[0] if v.dim() == 3 else 1
+            actual_H = self.num_heads
+            v_re = v_flat.view(actual_B, actual_H, HD)
+            U_re = self.U.view(actual_H, HD, R)
+            W_re = self.W.view(actual_H, HD, R)
+            
+            gamma_re = LowRankChristoffelFunction.apply(
+                v_re, U_re, W_re, self.clamp_val, self.enable_trace_normalization, True
+            )
+            gamma = gamma_re.view_as(v_flat)
+        else:
+            if v_flat.dim() == 3 and self.U.dim() == 3:
+                v_r = torch.einsum('bhd, hdr -> bhr', v_flat, self.U)
+                denom = 1.0 + torch.norm(v_r, dim=-1, keepdim=True)
+                phi = (v_r * v_r) / denom
+                gamma = torch.einsum('bhr, hdr -> bhd', phi, self.W)
+            else:
+                v_r = v_flat @ self.U
+                denom = 1.0 + torch.norm(v_r, dim=-1, keepdim=True)
+                phi = (v_r * v_r) / denom
+                W_t = self.W.transpose(-1, -2) if self.W.dim() == 3 else self.W.t()
+                gamma = phi @ W_t
+
+        x_in = self._get_features(x_flat)
+        mu = self.friction_gate(x_in, force=force)
+
+        if not use_cuda_fused:
+            gamma = self._normalize(gamma)
+            gamma = self.clamp_val * torch.tanh(gamma / self.clamp_val)
+
+        if B is not None:
+            gamma = gamma.view(original_shape)
+            mu = mu.view(original_shape)
+
+        # CONTRACT: Always return (gamma_pure, mu) for unified PhysicsEngine handling.
+        return gamma, mu

gfn/realizations/gssm/geometry/reactive.py

@@ -0,0 +1,109 @@
+"""
+ReactiveRiemannianGeometry — GFN V5
+Active-inference geometry: curvature reacts to system state.
+Migrated from gfn/geo/physical/reactive_field_geometry.py
+"""
+
+import torch
+import torch.nn as nn
+from typing import Optional, Union, Tuple, Dict, Any
+
+from ..constants import CURVATURE_CLAMP, EPS, DEFAULT_PLASTICITY, TOPOLOGY_TORUS
+from ..config.schema import PhysicsConfig
+from ..geometry.low_rank import LowRankRiemannianGeometry
+from ..registry import register_geometry
+
+# Default constants
+SINGULARITY_THRESHOLD = 0.5
+BLACK_HOLE_STRENGTH = 3.0
+SINGULARITY_GATE_SLOPE = 10.0
+
+
+@register_geometry('reactive')
+class ReactiveRiemannianGeometry(LowRankRiemannianGeometry):
+    """
+    Geometry that reacts to the system's own state via active inference.
+
+    Enhancements over LowRank:
+    1. Plasticity: Christoffel symbols scaled by kinetic energy (curv. amplification ≈ attention).
+    2. Singularities: Soft curvature amplification near semantic attractors.
+
+    Note: These are regularization/attention mechanisms, NOT physical manifold properties.
+    """
+
+    def __init__(self, dim: int, rank: int = 16, num_heads: int = 1,
+                 config: Optional[PhysicsConfig] = None):
+        super().__init__(dim, rank, num_heads=num_heads, config=config)
+        self.return_friction_separately = True
+
+        self.active_cfg = self.config.active_inference
+        self.plasticity = getattr(self.active_cfg, 'plasticity', DEFAULT_PLASTICITY)
+
+        sing_cfg = self.config.singularities
+        sing_enabled = getattr(sing_cfg, 'enabled', False)
+
+        if sing_enabled:
+            self.semantic_certainty_threshold = getattr(sing_cfg, 'threshold', SINGULARITY_THRESHOLD)
+            self.curvature_amplification_factor = getattr(sing_cfg, 'strength', BLACK_HOLE_STRENGTH)
+            gate_input_dim = dim * 2 if self.topology == TOPOLOGY_TORUS else dim
+            if num_heads > 1:
+                self.V_weight = nn.Parameter(torch.zeros(num_heads, gate_input_dim, 1))
+            else:
+                self.V = nn.Linear(gate_input_dim, 1)
+                nn.init.zeros_(self.V.weight)
+                nn.init.constant_(self.V.bias, -2.0)  # Start gate closed
+        else:
+            self.semantic_certainty_threshold = SINGULARITY_THRESHOLD
+            self.curvature_amplification_factor = BLACK_HOLE_STRENGTH
+            self.V = None
+
+    def _get_potential(self, x_in: torch.Tensor) -> Optional[torch.Tensor]:
+        """Compute singularity potential, returns None if disabled."""
+        if not getattr(self.config.singularities, 'enabled', False):
+            return None
+        if self.num_heads > 1:
+            return torch.sigmoid(torch.matmul(x_in.unsqueeze(-2), self.V_weight).squeeze(-2))
+        elif self.V is not None:
+            return torch.sigmoid(self.V(x_in))
+        return None
+
+    def forward(self, x: torch.Tensor, v: Optional[torch.Tensor] = None,
+                force: Optional[torch.Tensor] = None, **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        if v is None:
+            return torch.zeros_like(x)
+
+        # 1. Base curvature from LowRank
+        res = super().forward(x, v, force=force, **kwargs)
+        if isinstance(res, tuple):
+            gamma, mu = res
+        else:
+            gamma, mu = res, torch.zeros_like(v) if v is not None else torch.zeros_like(x)
+
+        if not self.active_cfg.enabled:
+            if self.return_friction_separately:
+                return gamma, mu
+            return gamma + mu * v if v is not None else gamma
+
+        # 2. Plasticity: scale curvature by kinetic energy
+        react_cfg = self.active_cfg.reactive_curvature
+        react_enabled = react_cfg.get('enabled', False) if isinstance(react_cfg, dict) else False
+        if react_enabled and self.plasticity > 0.0:
+            energy = torch.tanh(v.pow(2).mean(dim=-1, keepdim=True))
+            gamma = gamma * (1.0 + self.plasticity * energy)
+
+        # 3. Singularity amplification
+        if getattr(self.config.singularities, 'enabled', False) and x is not None:
+            x_in = torch.cat([torch.sin(x), torch.cos(x)], dim=-1) if self.topology == TOPOLOGY_TORUS else x
+            potential = self._get_potential(x_in)
+            if potential is not None:
+                gate_slope = getattr(self.config.singularities, 'gate_slope', SINGULARITY_GATE_SLOPE)
+                is_amplified = torch.sigmoid(gate_slope * (potential - self.semantic_certainty_threshold))
+                amp = 1.0 + is_amplified * (self.curvature_amplification_factor - 1.0)
+                gamma = gamma * amp
+                limit = self.curvature_amplification_factor * CURVATURE_CLAMP
+                gamma = limit * torch.tanh(gamma / limit)
+
+        if self.return_friction_separately:
+            return gamma, mu
+
+        return gamma + mu * v if v is not None else gamma

gfn/realizations/gssm/geometry/spherical.py

@@ -0,0 +1,47 @@
+import torch
+from typing import Optional, Union, Tuple, Any
+from ..geometry.base import BaseGeometry
+from ..registry import register_geometry
+from ..constants import EPS, TOPOLOGY_SPHERE
+
+@register_geometry(TOPOLOGY_SPHERE)
+class SphericalGeometry(BaseGeometry):
+    """
+    Spherical Geometry (Analytical).
+    Computes Christoffel symbols for a constant positive curvature space.
+    """
+    def __init__(self, dim: int, rank: int = 16, config: Optional[Any] = None, **kwargs):
+        super().__init__(config)
+        self.dim = dim
+
+    def christoffel_symbols(self, x: torch.Tensor) -> torch.Tensor:
+        """Analytical Christoffel symbols for S^n are typically zero in standard embedding or complex in other charts."""
+        # For our GFN purposes, we often use the simplified 'spherical_christoffel_torch' 
+        # which acts as a centering/restoring force towards the sphere surface.
+        return torch.zeros_like(x)
+
+    def forward(self, x: torch.Tensor, v: Optional[torch.Tensor] = None, 
+                force: Optional[torch.Tensor] = None, **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        if v is None:
+            return torch.zeros_like(x)
+            
+        # Simplified analytical spherical coupling (restoring force)
+        xv = torch.sum(x * v, dim=-1, keepdim=True)
+        vv = torch.sum(v * v, dim=-1, keepdim=True)
+        
+        # Gamma = -(2.0 * xv * v - vv * x)
+        gamma = -(2.0 * xv * v - vv * x)
+        
+        # Apply standard V5 clamping
+        clamp_val = getattr(self, 'clamp_val', 5.0)
+        return clamp_val * torch.tanh(gamma / clamp_val)
+
+    def metric_tensor(self, x: torch.Tensor) -> torch.Tensor:
+        """Identity metric for simplified spherical chart."""
+        return torch.ones_like(x)
+
+    def dist(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
+        """Great-circle distance approximation."""
+        dot = torch.sum(x1 * x2, dim=-1)
+        # Assuming points are on unit sphere
+        return torch.acos(torch.clamp(dot, -1.0 + EPS, 1.0 - EPS))

gfn/realizations/gssm/geometry/torus.py

@@ -0,0 +1,274 @@
+"""
+ToroidalRiemannianGeometry — GFN V5
+Full analytical torus geometry (canonical implementation).
+Replaces old stub. Migrated from gfn/geo/topological/toroidal_geometry.py
+"""
+
+import torch
+import torch.nn as nn
+from typing import Optional, Union, Tuple, Dict, Any
+
+from ..constants import (
+    EPS, TOPOLOGY_TORUS, CURVATURE_CLAMP,
+    DEFAULT_FRICTION, GATE_BIAS_CLOSED
+)
+from ..config.schema import PhysicsConfig
+from ..geometry.base import BaseGeometry
+from ..registry import register_geometry
+
+# Import modular friction component
+from ..physics.components.friction import FrictionGate
+
+# Torus-specific constants
+TOROIDAL_CURVATURE_SCALE = 0.1
+CLAMP_MIN_STRONG = 1e-4
+EPSILON_SMOOTH = 1e-9
+
+try:
+    import toroidal_cuda
+    HAS_CUDA_EXT = True
+except ImportError:
+    HAS_CUDA_EXT = False
+
+
+
+@register_geometry(TOPOLOGY_TORUS)
+class ToroidalRiemannianGeometry(BaseGeometry):
+    """
+    Curved torus geometry with exact (analytical) Christoffel symbols.
+
+    Metric (2D torus paired dimensions):
+      g_theta = r²
+      g_phi   = (R + r·cos θ)²
+
+    Generalizes to N dims by pairing (theta_0,phi_0), (theta_1,phi_1), ...
+    """
+
+    def __init__(self, dim: int = 64, rank: int = 16, num_heads: int = 1,
+                 config: Optional[PhysicsConfig] = None):
+        super().__init__(config)
+        self.dim = dim
+        self.num_heads = num_heads
+
+        topo = self.config.topology
+        
+        # CORRECCIÓN: Hacer R y r aprendibles según GFN_Paper_Complete.md Sección 5.1
+        # "where R_i are the (learnable) radii"
+        # Y 01_HYPER_TORUS.md Sección 2.2: R es radio mayor, r es radio menor
+        learnable_R = getattr(topo, 'learnable_R', True)
+        learnable_r = getattr(topo, 'learnable_r', True)
+        
+        if learnable_R:
+            # R como parámetro aprendible
+            self.R = nn.Parameter(torch.tensor(topo.R, dtype=torch.float32))
+        else:
+            # R como buffer no-entrenable
+            self.register_buffer('R', torch.tensor(topo.R, dtype=torch.float32))
+            
+        if learnable_r:
+            # r como parámetro aprendible
+            self.r = nn.Parameter(torch.tensor(topo.r, dtype=torch.float32))
+        else:
+            # r como buffer no-entrenable
+            self.register_buffer('r', torch.tensor(topo.r, dtype=torch.float32))
+        
+        self.topology = topo.type.lower()
+        self.clamp_val = self.config.stability.curvature_clamp
+
+        active_cfg = self.config.active_inference
+        self.active_cfg = active_cfg
+        rc = active_cfg.reactive_curvature
+        self.plasticity = rc.get('plasticity', getattr(active_cfg, 'plasticity', 0.0)) \
+            if isinstance(rc, dict) else getattr(active_cfg, 'plasticity', 0.0)
+
+        sing_cfg = self.config.singularities
+        self.singularity_threshold = sing_cfg.threshold
+        self.black_hole_strength = sing_cfg.strength
+
+        gate_input_dim = dim * 2  # [sin(x), cos(x)]
+        friction_mode = getattr(self.config.stability, 'friction_mode', 'static')
+        self.friction_gate = FrictionGate(
+            dim, gate_input_dim=gate_input_dim, mode=friction_mode, num_heads=num_heads
+        )
+
+        # Optional singularity potential gate
+        if getattr(sing_cfg, 'enabled', False):
+            if num_heads > 1:
+                self.V_weight = nn.Parameter(torch.zeros(num_heads, gate_input_dim, 1))
+                self.V_bias = nn.Parameter(torch.full((num_heads, 1), GATE_BIAS_CLOSED))
+            else:
+                self.V = nn.Linear(gate_input_dim, 1)
+                nn.init.zeros_(self.V.weight)
+                nn.init.constant_(self.V.bias, GATE_BIAS_CLOSED)
+        else:
+            self.V = None
+
+    def validate_dimensions(self, x: torch.Tensor):
+        if x.shape[-1] % 2 != 0:
+            raise ValueError(
+                f"ToroidalGeometry requires even dim for (θ,φ) pairing. Got {x.shape[-1]}"
+            )
+
+    def metric_tensor(self, x: torch.Tensor) -> torch.Tensor:
+        self.validate_dimensions(x)
+        g = torch.ones_like(x)
+        for i in range(0, x.shape[-1], 2):
+            th = x[..., i]
+            g[..., i] = self.r ** 2
+            g[..., i + 1] = (self.R + self.r * torch.cos(th)) ** 2
+        return g
+
+    def connection(self, v: torch.Tensor, w: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+        d = x.shape[-1]
+        gamma = torch.zeros_like(v)
+        for i in range(0, (d // 2) * 2, 2):
+            th = x[..., i]
+            v_th, v_ph = v[..., i], v[..., i + 1]
+            w_th, w_ph = w[..., i], w[..., i + 1]
+            denom = torch.clamp(self.R + self.r * torch.cos(th), min=CLAMP_MIN_STRONG)
+            term_th = (denom * torch.sin(th) / (self.r + EPSILON_SMOOTH)) * TOROIDAL_CURVATURE_SCALE
+            gamma[..., i] = term_th * (v_ph * w_ph)
+            term_ph = (-(self.r * torch.sin(th)) / (denom + EPSILON_SMOOTH)) * TOROIDAL_CURVATURE_SCALE
+            gamma[..., i + 1] = term_ph * (v_ph * w_th + v_th * w_ph)
+        return gamma
+
+    def forward(self, x: torch.Tensor, v: Optional[torch.Tensor] = None,
+                force: Optional[torch.Tensor] = None, **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        if v is None:
+            return torch.zeros_like(x)
+
+        self.validate_dimensions(x)
+
+        if HAS_CUDA_EXT and x.is_cuda and v is not None and v.is_cuda:
+            gamma = toroidal_cuda.forward(
+                x, v, self.R, self.r, 
+                TOROIDAL_CURVATURE_SCALE, EPSILON_SMOOTH, CLAMP_MIN_STRONG
+            )
+        else:
+            is_odd = (self.dim % 2 != 0)
+            x_pad = torch.nn.functional.pad(x, (0, 1)) if is_odd else x
+            v_pad = torch.nn.functional.pad(v, (0, 1)) if is_odd else v
+    
+            th = x_pad[..., 0::2]
+            v_th = v_pad[..., 0::2]
+            v_ph = v_pad[..., 1::2]
+    
+            denom = torch.clamp(self.R + self.r * torch.cos(th), min=CLAMP_MIN_STRONG)
+            term_th = (denom * torch.sin(th) / (self.r + EPSILON_SMOOTH))
+            term_ph = -(self.r * torch.sin(th)) / (denom + EPSILON_SMOOTH)
+    
+            gamma_th = term_th * (v_ph ** 2) * TOROIDAL_CURVATURE_SCALE
+            gamma_ph = 2.0 * term_ph * v_ph * v_th * TOROIDAL_CURVATURE_SCALE
+    
+            half = x.shape[-1] // 2
+            gamma = torch.zeros_like(x)
+            gamma[..., 0::2] = gamma_th[..., :half + x.shape[-1] % 2]
+            gamma[..., 1::2] = gamma_ph[..., :half]
+
+        # Friction gate
+        x_in = torch.cat([torch.sin(x), torch.cos(x)], dim=-1)
+        mu = self.friction_gate(x_in, force=force)
+        if mu.shape[-1] != v.shape[-1]:
+            if mu.shape[-1] == 2 * v.shape[-1]:
+                mu = (mu[..., :v.shape[-1]] + mu[..., v.shape[-1]:]) / 2.0
+
+        # Active inference
+        if self.active_cfg.enabled:
+            rc = self.active_cfg.reactive_curvature
+            react_enabled = rc.get('enabled', False) if isinstance(rc, dict) else False
+            if react_enabled and self.plasticity > 0.0:
+                energy = torch.tanh(v.pow(2).mean(dim=-1, keepdim=True))
+                gamma = gamma * (1.0 + self.plasticity * energy)
+
+            if getattr(self.config.singularities, 'enabled', False):
+                if self.num_heads > 1:
+                    potential = torch.sigmoid(
+                        torch.matmul(x_in.unsqueeze(-2), self.V_weight).squeeze(-2)
+                        + self.V_bias
+                    )
+                elif self.V is not None:
+                    potential = torch.sigmoid(self.V(x_in))
+                else:
+                    potential = None
+
+                if potential is not None:
+                    soft_m = torch.sigmoid(5.0 * (potential - self.singularity_threshold))
+                    gamma = gamma * (1.0 + soft_m * (self.black_hole_strength - 1.0))
+
+        # CONTRACT: Always return (gamma_pure, mu) — engine applies friction, not geometry.
+        # This unifies the contract with LowRankRiemannianGeometry (P0.2 fix).
+        return gamma, mu
+
+    def project(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.atan2(torch.sin(x), torch.cos(x))
+
+    def dist(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
+        diff = x1 - x2
+        return torch.norm(torch.atan2(torch.sin(diff), torch.cos(diff)), dim=-1)
+
+
+@register_geometry('flat_torus')
+class FlatToroidalRiemannianGeometry(BaseGeometry):
+    def __init__(self, dim: int = 64, rank: int = 16, num_heads: int = 1,
+                 config: Optional[PhysicsConfig] = None):
+        super().__init__(config)
+        self.dim = dim
+        self.num_heads = num_heads
+        topo = self.config.topology
+        
+        # CORRECCIÓN: Hacer R y r aprendibles también en FlatTorus
+        learnable_R = getattr(topo, 'learnable_R', True)
+        learnable_r = getattr(topo, 'learnable_r', True)
+        
+        if learnable_R:
+            self.R = nn.Parameter(torch.tensor(topo.R, dtype=torch.float32))
+        else:
+            self.register_buffer('R', torch.tensor(topo.R, dtype=torch.float32))
+            
+        if learnable_r:
+            self.r = nn.Parameter(torch.tensor(topo.r, dtype=torch.float32))
+        else:
+            self.register_buffer('r', torch.tensor(topo.r, dtype=torch.float32))
+            
+        self.topology = topo.type.lower()
+        self.return_friction_separately = True
+        gate_input_dim = dim * 2
+        friction_mode = getattr(self.config.stability, 'friction_mode', 'static')
+        self.friction_gate = FrictionGate(
+            dim, gate_input_dim=gate_input_dim, mode=friction_mode, num_heads=num_heads
+        )
+
+    def validate_dimensions(self, x: torch.Tensor):
+        if x.shape[-1] % 2 != 0:
+            raise ValueError(
+                f"FlatToroidalGeometry requires even dim for (θ,φ) pairing. Got {x.shape[-1]}"
+            )
+
+    def metric_tensor(self, x: torch.Tensor) -> torch.Tensor:
+        self.validate_dimensions(x)
+        return torch.ones_like(x)
+
+    def forward(self, x: torch.Tensor, v: Optional[torch.Tensor] = None,
+                force: Optional[torch.Tensor] = None, **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        if v is None:
+            return torch.zeros_like(x)
+
+        self.validate_dimensions(x)
+        x_in = torch.cat([torch.sin(x), torch.cos(x)], dim=-1)
+        mu = self.friction_gate(x_in, force=force)
+        if mu.shape[-1] != v.shape[-1]:
+            if mu.shape[-1] == 2 * v.shape[-1]:
+                mu = (mu[..., :v.shape[-1]] + mu[..., v.shape[-1]:]) / 2.0
+            else:
+                mu = mu.expand_as(v)
+        gamma = torch.zeros_like(v)
+        if self.return_friction_separately:
+            return gamma, mu
+        return gamma + mu * v
+
+    def project(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.atan2(torch.sin(x), torch.cos(x))
+
+    def dist(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
+        diff = x1 - x2
+        return torch.norm(torch.atan2(torch.sin(diff), torch.cos(diff)), dim=-1)