OjasX v0.6.0 — The Open-Source CUDA.

A new Mathematical & MLIR Compiler for Training PyTorch Models on Any Hardware

🌟 Project Motto

"Democratizing Deep Learning Compute: Breaking the Hardware Monopoly through Pure Mathematics,and Open Compiler Infrastructure."

---

⚔️ Breaking the Computing Monopoly & Proprietary CUDA Lock-In

Today, frontier AI research and large-scale model training are artificially gatekept behind a wall of proprietary vendor lock-in. Closed proprietary ecosystems—most notably CUDA and proprietary accelerator fabrics—dictate who gets to train modern neural networks. This compute cartel forces universities, independent researchers, and startups to pay tens of thousands of dollars for enterprise-tier accelerators or submit to cloud compute monopolies simply to avoid Out-Of-Memory (OOM) crashes during backpropagation.

OjasX was built from first principles to shatter this monopoly.

By reformulating backpropagation through Category-Theoretic Dagger Monoidal Categories, volume-preserving Hamiltonian Symplectic flows, and a 3-dialect MLIR compiler pipeline, OjasX makes hardware specifics secondary to mathematical invariants:

• 0 Bytes DRAM for Attention Scores: Attention scores are tiled directly in fast hardware registers and on-chip SRAM via online-softmax pullbacks. The quadratic [B, H, S, S] tensor is never written to DRAM.
• 0 Bytes DRAM for Vocabulary Logits: Large language model output heads stream vocabulary projections tile-by-tile, eliminating multi-gigabyte [B, S, V] logit scratchpads.
• 0 Bytes Persistent Gradients: The HolonomicGradFreeOptimizer absorbs parameter updates layer-by-layer during the backward pass and wipes .grad buffers immediately, freeing up gigabytes of memory.
• O(1) Depth-Independent Reversible Memory: Symplectic coupling chains reconstruct intermediate states backwards on the fly, allowing 64-layer to 128-layer deep backbones to run in the exact same activation footprint as 1 layer.
• True Multi-Vendor Hardware Freedom: OjasX compiles and executes natively on NVIDIA GPUs (CUDA PTX), AMD GPUs (ROCm/HIP), Intel GPUs (SPIR-V/SYCL), Apple Silicon (Metal 4 MSL / AMX), and standard x86/ARM CPUs.

High-performance AI training belongs on every developer's desk—not behind a proprietary vendor tollgate.

---

Quick Navigation & Core Capabilities

---

Step-by-Step Installation & Setup

Step 1: Install OjasX (Builds Native C++17 / GPU Kernels Automatically)

git clone https://github.com/mkrishna793/OjasX-open-source-CUDA.git
cd OjasX-open-source-CUDA

# Install with native C++17 / SIMD / Metal / GPU kernel compilation:
pip install -e . --no-build-isolation

Note: If a C++ compiler is not present on your system, OjasX automatically falls back to its 100% compatible PyTorch operator backend with zero manual configuration.

Step 2: Verify Your Hardware & Installation

# Check detected CPU/GPU silicon profile, cache sizes, and native kernel status:
ojasx-info

# Run live mathematical & gradient parity diagnostics across 5 model families:
ojasx-doctor

---

30-Second Universal Quickstart (Works on Any Model & Any Hardware)

import torch
import torch.nn as nn
import ojasx2
from ojasx2.grad_free_optimizer import HolonomicGradFreeOptimizer

# Step 1: Define or load any standard PyTorch model (LLM, ViT, CNN, or MLP)
class MyPyTorchModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(128, 256), nn.GELU(), nn.Linear(256, 128))

    def forward(self, x):
        return self.net(x)

device = (
    "cuda" if torch.cuda.is_available()
    else "mps" if (hasattr(torch.backends, "mps") and torch.backends.mps.is_available())
    else "xpu" if (hasattr(torch, "xpu") and torch.xpu.is_available())
    else "cpu"
)
model = MyPyTorchModel().to(device)

# Step 2: Optimize model layers AND lower through the 3-dialect MLIR pipeline in 1 line
model, report = ojasx2.optimize_model(
    model,
    batch_tokens=512,
    prefer_unified_mlp=True,
    compile_mlir=True,
    mlir_target="auto",
)
print(report.format_table())
print("MLIR Pipeline Summary:", report.mlir_pipeline)

# Step 3: Attach the Zero-Persistent-Grad Optimizer (0 bytes of persistent .grad memory)
optimizer = HolonomicGradFreeOptimizer(
    model.parameters(), lr=3e-4, algorithm="adamw", weight_decay=0.01, max_grad_norm=1.0
)

# Step 4: Train normally — weights update in-place during loss.backward()!
model.train()
inputs = torch.randn(8, 128, device=device)
targets = torch.randn(8, 128, device=device)
loss_fn = nn.MSELoss()

optimizer.zero_grad()
loss = loss_fn(model(inputs), targets)
loss.backward()      # Applies parameter updates layer-by-layer & frees .grad immediately
optimizer.step()

---

Supported Hardware & Automatic Target Routing

OjasX runs natively across all major GPU and CPU platforms without code changes:

---

Copy-Paste Recipes for 7 Common Use Cases

Use Case 1: Training Any Custom Transformer or LLM (With Fused Zero-Logit Cross-Entropy)

When your model has an output vocabulary projection (self.lm_head), ojasx2.optimize_model automatically upgrades it to HolonomicFusedLinearCrossEntropyHead:

• Calling model.lm_head(hidden, targets) computes the cross-entropy loss directly in streaming tiles (0 bytes of [B, S, V] logits in VRAM).
• Calling model.lm_head(hidden) (without targets, e.g., during inference/generation) returns standard [B, S, V] logits.

import torch
import torch.nn as nn
import torch.nn.functional as F
import ojasx2
from ojasx2.grad_free_optimizer import HolonomicGradFreeOptimizer

class CustomTransformerLM(nn.Module):
    def __init__(self, vocab_size=32000, d_model=512, nhead=8, num_layers=6):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        self.layers = nn.ModuleList([
            nn.TransformerEncoderLayer(
                d_model=d_model, nhead=nhead, dim_feedforward=d_model * 4,
                activation="gelu", batch_first=True, norm_first=True
            )
            for _ in range(num_layers)
        ])
        self.norm = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, input_ids, targets=None):
        h = self.embed(input_ids)
        for layer in self.layers:
            h = layer(h)
        h = self.norm(h)
        if targets is not None:
            # Uses OjasX streaming cross-entropy (0 bytes full-vocab logits in DRAM)
            return self.lm_head(h, targets)
        return self.lm_head(h)

device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
model = CustomTransformerLM().to(device)

# Optimize Attention, LayerNorm, GELU, Linear, and LM Head:
model, report = ojasx2.optimize_model(model, batch_tokens=256, compile_mlir=True)
print("Applied Optimizations:", report.action_counts)

optimizer = HolonomicGradFreeOptimizer(model.parameters(), lr=1e-3, algorithm="adamw")

# Training step
x = torch.randint(0, 32000, (4, 64), device=device)
y = torch.randint(0, 32000, (4, 64), device=device)
loss = model(x, targets=y)
loss.backward()
optimizer.step()
optimizer.zero_grad()

# Inference / Eval step
model.eval()
with torch.no_grad():
    logits = model(x)  # Shape: [4, 64, 32000]

---

Use Case 2: Fine-Tuning Hugging Face Models & LoRA / Frozen Backbones

When you freeze base weights (param.requires_grad = False), ojasx2.optimize_model automatically detects the LORA_FROZEN_BACKBONE profile and wraps frozen nn.Linear and nn.Conv* layers in FrozenLinearZeroTape / FrozenConvZeroTape, dropping 100% of input activation storage on frozen layers.

import torch
from transformers import AutoModelForCausalLM
import ojasx2
from ojasx2.grad_free_optimizer import HolonomicGradFreeOptimizer

device = "cuda" if torch.cuda.is_available() else "cpu"
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B", torch_dtype=torch.bfloat16).to(device)

# Optional: Freeze base backbone and train only selected layers or LoRA adapters
for name, param in model.named_parameters():
    if "lora" not in name and "lm_head" not in name:
        param.requires_grad = False

# Optimize model: rewrites Attention, RMSNorm, SwiGLU MLP, and skips tape on frozen weights
model, report = ojasx2.optimize_model(model, prefer_unified_mlp=True)
print(report.format_table())

# Train only trainable parameters with 0 persistent .grad buffers
trainable_params = [p for p in model.parameters() if p.requires_grad]
optimizer = HolonomicGradFreeOptimizer(trainable_params, lr=2e-4, algorithm="adamw")

input_ids = torch.randint(0, 32000, (2, 512), device=device)
outputs = model(input_ids, labels=input_ids)
outputs.loss.backward()
optimizer.step()
optimizer.zero_grad()

---

Use Case 3: Training Ultra-Deep Models (32–128+ Layers) in O(1) Constant Activation Memory

Use SymplecticTransformerBlock or symplectic_coupling_chain to train arbitrarily deep networks where activation memory stays strictly constant ($\mathcal{O}(1)$) with respect to depth $D$. During backward, states are reconstructed on the fly via exact Hamiltonian inverse steps: $$q_2^{(i)} = q_2^{(i+1)} - g_i(q_1^{(i+1)}), \qquad q_1^{(i)} = q_1^{(i+1)} - f_i(q_2^{(i)})$$

import torch
import torch.nn as nn
from ojasx2 import (
    SymplecticTransformerBlock,
    symplectic_coupling_chain,
    audit_symplectic_reversibility,
    HolonomicGradFreeOptimizer,
)

class DeepReversibleBackbone(nn.Module):
    def __init__(self, dim=512, depth=64, num_heads=8, vocab_size=32000):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, dim)
        # 64 layers with the exact same activation tape footprint as 1 layer!
        self.blocks = nn.ModuleList([
            SymplecticTransformerBlock(dim=dim, num_heads=num_heads, split_dim=True)
            for _ in range(depth)
        ])
        self.head = nn.Linear(dim, vocab_size, bias=False)

    def forward(self, input_ids):
        x = self.embed(input_ids)
        half = x.size(-1) // 2
        # Execute entire 64-layer chain saving ONLY the final boundary state (y1, y2)
        y1, y2 = symplectic_coupling_chain(x[..., :half], x[..., half:], self.blocks)
        return self.head(torch.cat([y1, y2], dim=-1))

model = DeepReversibleBackbone(dim=256, depth=32)
optimizer = HolonomicGradFreeOptimizer(model.parameters(), lr=1e-3, algorithm="adamw")

tokens = torch.randint(0, 32000, (2, 64))
logits = model(tokens)
logits.sum().backward()
optimizer.step()

# Verify mathematical invertibility of the chain (Error <= 4.44e-16 in FP64)
audit = audit_symplectic_reversibility(torch.randn(2, 64, 256), model.blocks)
print("Reversible:", audit["is_reversible"], "| Max Reconstruction Error:", audit["max_error"])

---

Use Case 4: Using Modular Low-Memory Primitives Directly in Custom Architectures

You can import OjasX's individual layers and functional operators directly into any custom model:

import torch
from ojasx2 import (
    HolonomicDaggerAttention,
    holonomic_dagger_attention,
    HolonomicUnifiedTransformerMLP,
    HolonomicFusedSwiGLUMLP,
    HolonomicNativeRMSNorm,
    HolonomicLayerNorm,
    HolonomicFusedLinearCrossEntropy,
)

# 1. Tiled Online-Softmax Attention (supports GQA, causal masks, boolean/additive attn_mask, key_padding_mask)
q = torch.randn(2, 8, 512, 64, requires_grad=True)   # 8 Query heads
k = torch.randn(2, 2, 512, 64, requires_grad=True)   # 2 KV heads (4x GQA)
v = torch.randn(2, 2, 512, 64, requires_grad=True)
pad_mask = torch.zeros(2, 512, dtype=torch.bool)     # Optional [B, S_kv] padding mask
attn_out = holonomic_dagger_attention(q, k, v, is_causal=True, key_padding_mask=pad_mask)
attn_out.sum().backward()

# 2. Single-Node Unified Transformer MLP (RMSNorm -> SwiGLU -> DownProj with 0 intermediate tape)
unified_mlp = HolonomicUnifiedTransformerMLP(dim=512, intermediate_dim=1536)
x_mlp = torch.randn(2, 128, 512, requires_grad=True)
y_mlp = unified_mlp(x_mlp)
y_mlp.sum().backward()

# 3. Streaming Tiled Cross-Entropy Loss (0 bytes of [B*S, V] logits in DRAM)
fused_ce = HolonomicFusedLinearCrossEntropy(in_features=512, vocab_size=128000, ignore_index=-100)
hidden = torch.randn(256, 512, requires_grad=True)
targets = torch.randint(0, 128000, (256,))
targets[:64] = -100  # Padded tokens are automatically culled before backward GEMMs!
loss = fused_ce(hidden, targets)
loss.backward()

---

Use Case 5: Computer Vision (CNNs / ResNets / UNets) & Tabular MLPs

OjasX automatically compresses ReLU activations by 32× using hardware 1-bit ballot masks (PackedReLU / HolonomicBitPackedLinear), replaces LayerNorm/GroupNorm with compact scalar-statistics lenses, and eliminates activation storage on frozen convolutional stems.

import torch
import torch.nn as nn
import ojasx2

class VisionBackboneClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.stem = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.stem.weight.requires_grad_(False)  # Frozen stem -> FrozenConvZeroTape (0 input tape!)
        self.norm = nn.GroupNorm(8, 64)         # -> HolonomicGroupNorm (compact group stats)
        self.act = nn.ReLU()                    # -> PackedReLU (1-bit sign mask: 32x smaller than FP32)
        self.head = nn.Sequential(
            nn.Linear(64, 128),
            nn.ReLU(),
            nn.Linear(128, 10),
        )

    def forward(self, x):
        h = self.act(self.norm(self.stem(x)))
        return self.head(h.mean(dim=(-1, -2)))

model = VisionBackboneClassifier()
# Set prefer_1bit_relu2_for_relu=True to fuse Linear + ReLU into 1-bit ballot C++/GPU kernels:
model, report = ojasx2.optimize_model(model, prefer_1bit_relu2_for_relu=True)
print(report.format_table())

---

Use Case 6: 3-Dialect MLIR Compilation & torch.compile Hardware Codegen

Lower any PyTorch model through OjasX's 3-dialect MLIR pipeline (ojasx.holonomic $\to$ ojasx.subgraph $\to$ ojasx.tile) or use torch.compile(..., backend="ojasx_mlir"):

import torch
import torch.nn as nn
import ojasx2

model = nn.Sequential(
    nn.Linear(128, 256),
    nn.ReLU(),
    nn.Linear(256, 128),
)

# Option A: Compile into an executable MLIR module and inspect all 3 IR dialects + hardware code
exec_mod = ojasx2.compile_mlir_executable(model, target="auto")
print("SSA Verified:", exec_mod.mlir_artifacts.ssa_verified)
print(exec_mod.holonomic_ir)   # High-level ojasx.holonomic dialect
print(exec_mod.subgraph_ir)    # Memory-planned ojasx.subgraph dialect
print(exec_mod.tile_ir)        # Hardware-tiled ojasx.tile dialect
print(exec_mod.target_code)    # Emitted NVVM PTX / ROCDL / SPIR-V / Metal / C++ code

# Option B: Emit source directly for any target ("nvvm", "rocdl", "spirv", "metal", "tilelang", "cpp")
ptx_source = ojasx2.compile_model_to_target(model, target="nvvm")
metal_source = ojasx2.compile_model_to_target(model, target="metal")

# Option C: Use as a native TorchDynamo backend with torch.compile
compiled_model = torch.compile(model, backend="ojasx_mlir")
x = torch.randn(16, 128, requires_grad=True)
out = compiled_model(x)
out.sum().backward()

---

Use Case 7: Multi-GPU / Multi-Node Pipeline Training (HolonomicVirtualMeshEngine)

When training 13B–70B+ models whose weights exceed a single GPU's memory, partition stages across multiple GPUs or machines over standard TCP/IP with automatic 8-bit block-scaled boundary compression (4× lower network bandwidth):

import torch
import torch.nn as nn
from ojasx2.mesh_hal import HolonomicVirtualMeshEngine, MeshNodeSpec

stages = [
    nn.Sequential(nn.Linear(512, 512), nn.SiLU()),
    nn.Sequential(nn.Linear(512, 512), nn.SiLU()),
]
nodes = [
    MeshNodeSpec(node_id=0, host="127.0.0.1", port=29500, device="cpu", is_local=True),
    MeshNodeSpec(node_id=1, host="127.0.0.1", port=29501, device="cpu", is_local=True),
]

mesh = HolonomicVirtualMeshEngine(stages=stages, nodes=nodes, compress_wire=True)
x = torch.randn(8, 512, requires_grad=True)
out = mesh.forward_backward(x, loss_fn=lambda y: y.pow(2).mean())
print("Wire compression ratio:", mesh.stats.compression_ratio)

---

Quick Parameter & Tuning Guidelines

---

Officially Verified Benchmarks (PyTorch Eager vs. OjasX v0.6.0)

All metrics below are measured and recorded in the verified test database matrix reports/OJASX_V05_TEST_DATABASE_MATRIX.json using Meta's BackendBench testing harness:

A. The 1.04 Billion Parameter Transformer LLM Benchmark

(18-Layer LLaMA/Qwen-style Transformer: Dim=2048, Hidden=5504, 16 Attention Heads, Vocab=32,000, Sequence Length=512, Full Precision FP32)

💡 The 6.35 GB Reality: How a 1B+ Model Trains on an Everyday Laptop

When training the 1.04 Billion Parameter LLM in full 32-bit floating-point precision (FP32), here is where every byte of RAM actually goes:

┌────────────────────────────────────────────────────────────────────────┐
│ 1.04B PARAMETER MODEL TRAINING MEMORY (FULL PRECISION FP32)            │
├────────────────────────────────────────────────────────────────────────┤
│ • Physical Model Weights (1.04B params in FP32)  :  4.16 GB            │
│   (Raw uncompressed parameter tensors in memory)                       │
│                                                                        │
│ • OjasX Activation Scratchpad                    :  0.69 GB            │
│   (Reduced from 2.48 GB in PyTorch via tiled online pullbacks)         │
│                                                                        │
│ • Persistent Gradients (.grad buffers)           :  0.00 GB            │
│   (Freed immediately layer-by-layer; PyTorch holds 4.16 GB)            │
│                                                                        │
│ • Optimizer Working Buffers                      : ~1.50 GB            │
│   (In-hook updates; standard PyTorch AdamW holds 8.32 GB)              │
├────────────────────────────────────────────────────────────────────────┤
│ 👉 TOTAL PYTHON PROCESS MEMORY FOR OJASX         :  6.35 GB            │
│    (Standard PyTorch eager alone demands         : 19.12 GB)           │
└────────────────────────────────────────────────────────────────────────┘

• Standard PyTorch Eager: Requires ~19.12 GB just for the Python process. On a 16 GB or 24 GB laptop with Chrome, your IDE, and macOS running (~8 GB), standard PyTorch immediately crashes with an Out Of Memory (OOM) panic.
• OjasX v0.6.0: The entire training process takes only 6.35 GB. Even with Google Chrome, multiple browser tabs, your IDE, and macOS running simultaneously, total machine memory stays at ~14.5 GB (comfortable Green memory pressure on a 24 GB Mac) with over 9.5 GB of free headroom!

B. Comprehensive Benchmark Across All 6 Workload Families

---

Visual Comparison Charts

A. 1B Model Training Step Latency (Lower is Better)

Official PyTorch:  ████████████████████████████████████████ (101.4 sec)
OjasX v0.6.0:      ███ (7.1 sec)   [⚡ 14.24× FASTER SPEEDUP]

B. Whole-Model 1B LLM Activation Memory per Step (Lower is Better)

Official PyTorch:  ████████████████████████████████████████ (2,482 MB)
OjasX v0.6.0:      ███████████ (692 MB)   [-72.09% REDUCTION]

C. Intermediate Scratchpad Memory Materialized in DRAM

Attention Scores [B, H, S, S]:
  PyTorch:  ████████████████████ (1,200 MB)
  OjasX:    [ 0 Bytes ]

Vocabulary Logits [B, S, V]:
  PyTorch:  █████████ (524 MB)
  OjasX:    [ 0 Bytes ]

SwiGLU MLP Intermediates:
  PyTorch:  ████████████ (758 MB)
  OjasX:    ████ (262 MB) [-69.2%]

D. Memory Savings vs Sequence Length Scaling

xychart-beta
    title "Peak Activation Memory vs Sequence Length (1B Parameter LLM)"
    x-axis [512, 1024, 2048, 4096]
    y-axis "Activation VRAM (MB)" 0 --> 16000
    bar [2482, 5100, 10400, 16200]
    bar [692, 1420, 2900, 4800]

E. Run the Auto-Hardware Benchmark on Your Machine (1-Command)

OjasX includes an automatic hardware- and RAM-aware benchmarking utility that detects your active accelerator (NVIDIA CUDA, AMD ROCm, Apple Metal MPS, Intel XPU, or CPU), measures available memory, selects the largest fitting Transformer LLM tier (125M up to 13B+), and records peak activation tape memory, 0-byte persistent .grad consumption, and forward+backward step runtime:

# Auto-detect hardware accelerator, RAM, and largest fitting LLM tier:
python tools/auto_hardware_benchmark.py

# Or explicitly customize model tier, precision, and training mode:
python tools/auto_hardware_benchmark.py --tier 1b --dtype fp32 --optimizer sgd
python tools/auto_hardware_benchmark.py --tier 7b --dtype bf16 --optimizer adamw --lora

---

Project Structure for Developers & Researchers

The OjasX codebase is organized into clean, mathematically rigorous subsystems designed for rapid research experimentation and production hardware compiler lowering:

OjasX-open-source-CUDA/
├── ojasx2/                         # Core Python Library & Autograd Subsystems
│   ├── __init__.py                 # Top-level API exports (v0.6.0)
│   ├── smart_router.py             # Automatic model inspector & layer rewriter (optimize_model)
│   ├── dagger_attention.py         # In-register online-softmax attention (HolonomicDaggerAttention)
│   ├── fused_cross_entropy.py      # Streaming token-vocab cross-entropy (HolonomicFusedLinearCrossEntropy)
│   ├── grad_free_optimizer.py      # In-hook zero-persistent-grad optimizer (HolonomicGradFreeOptimizer)
│   ├── symplectic_engine.py        # Hamiltonian volume-preserving reversible stack (SymplecticTransformerBlock)
│   ├── mlir_dialects.py            # 3-Dialect MLIR pipeline (holonomic -> subgraph -> tile) & codegen
│   ├── fx_subgraph_compiler.py     # PyTorch FX subgraph pattern matcher & operator fusions
│   ├── tile_hal.py                 # TileLang hardware abstraction layer, schedules & Z3 provers
│   ├── tile_fallback.py            # Portable pure-PyTorch reference kernels across all backends
│   ├── holonomic_layers.py         # HolonomicLayerNorm, GroupNorm, PackedReLU, FrozenConvZeroTape
│   ├── mesh_hal.py                 # Distributed TCP virtual mesh engine with 8-bit wire compression
│   ├── aot_holonomic_compiler.py   # torch.compile AOTAutograd backend integration
│   ├── doctor.py                   # Diagnostic CLI (ojasx-doctor) for live system verification
│   └── csrc/                       # Native C++17, Accelerate AMX, and Metal Hardware Kernels
│       ├── ojasx_tile_kernels.cpp  # CPU AVX2/NEON/AMX GEMM, 1-bit ballot ReLU2, L1 RMSNorm
│       ├── ojasx_tilelang_metal.mm # Metal 4 MSL JIT compiler & universal codegen (CUDA/ROCm/SPIR-V)
│       └── ojasx_openreg_guard.cpp # PrivateUse1 (torch.ojasx) device guard implementation
├── lean/                           # Lean 4 Formal Mathematical Verification
│   └── OjasXMath/                  # 33 Lean 4 files: 558 formally proven theorems (0 sorry)
│       ├── AbstractEngine.lean     # Categorical reverse derivative foundation
│       ├── SymplecticCotangent.lean# Formal proof of unit Jacobian det(J)=1 and volume preservation
│       └── StreamedDagger.lean     # Dagger monoidal contraction theorems
├── rocq/                           # Rocq 9.2 (Coq) Constructive Formal Proofs
│   ├── *.v                         # 25 Gallina verification specifications (0 Admitted, 0 Axiom)
│   └── extracted/*.ml              # Extracted OCaml verified reference kernel implementations
├── tests/                          # 840+ Comprehensive Automated Test Suites
│   ├── test_challenger_stress_v05.py # Stress tests (deep Symplectic chains, extreme lengths, prime dims)
│   ├── test_opcheck_compliance.py    # PyTorch contract compliance (schema, autograd, FakeTensor, AOT)
│   ├── test_opinfo_compliance.py     # 249-combination OpInfo multi-stride and multi-dtype matrix
│   ├── test_backendbench_matrix.py   # Meta BackendBench real-world model suites (including 1.04B LLM)
│   └── test_review_regressions.py    # Regression tests for mathematical pullbacks and edge cases
├── tools/                          # Developer & Verification CLI Utilities
│   ├── auto_hardware_benchmark.py  # Automatic hardware & RAM-aware community benchmark runner
│   ├── test_universal_recipe.py    # End-to-end multi-device universal recipe verification script
│   ├── verify_static_and_symbolic.py # 8-gate static, symbolic, MLIR SSA, and formal proof validator
│   └── self_audit_complete.py      # Complete AST symbol resolver and cross-module auditor
└── reports/                        # Machine-Readable Benchmark & Audit Databases
    └── OJASX_V05_TEST_DATABASE_MATRIX.json # Authoritative test database matrix

For Researchers: Adding Custom Mathematical Lenses

Researchers can define custom forward-backward morphisms without writing manual backward kernels. Define a custom activation or contraction lens in ojasx2/lenses.py or ojasx2/holonomic_layers.py:

1. Subclass ojasx2.Lens or torch.autograd.Function.
2. Register your pullback via contravariant dagger adjoints $(g \circ f)^\dagger = f^\dagger \circ g^\dagger$.
3. Plug into SmartWorkloadRouter by registering your custom replacement in _rewrite_children.

For Compiler Engineers: Adding New Hardware Targets

OjasX's 3-dialect MLIR pipeline makes hardware backends modular:

1. ojasx.holonomic: Captures high-level algebraic morphisms and operator fusions.
2. ojasx.subgraph: Performs contraction reassociation and memory planning.
3. ojasx.tile: Emits explicit tile allocations (T.alloc_shared, T.alloc_fragment) and schedules.
4. To add a new hardware emitter (e.g. TPU, Tenstorrent, or custom ASIC), implement an emitter function in ojasx2/mlir_dialects.py targeting your ISA or intermediate representation.

---

Official PyTorch Contract Enforcement & Verification

To guarantee that OjasX is production-grade and 100% compliant with PyTorch internals, every operator is verified through PyTorch's official validation suites with zero synthetic or fake tests:

• torch.library.opcheck: 100% pass rate across:
  • test_schema: C++ operator argument and return type contracts.
  • test_autograd_registration: Exact gradient formula and backward registration parity.
  • test_faketensor: Symbolic shape propagation and meta-device compatibility.
  • test_aot_dispatch_dynamic: Dynamic sequence length tracing under torch.compile.
• Metadata-Driven OpInfo Matrix: 249 parameter combinations passed with 0 failures across CPU and Apple Silicon GPU (mps), dtypes (float32, float16, bfloat16), and non-contiguous memory layouts (transposed, sliced, permuted).
• Full Project Test Suite: 840 tests verified: 839 passed, 0 failures, 1 hardware-dependent skip across Meta BackendBench, OpCheck, OpInfo, and MLIR compiler stress tests.

---

Verification & Testing Commands

# 1. Run the universal training recipe verification script (CPU / GPU / MPS):
python3 tools/test_universal_recipe.py

# 2. Run the static, symbolic, MLIR SSA, and formal proof audit suite:
python3 tools/verify_static_and_symbolic.py

# 3. Run the full pytest suite (840 tests across OpCheck, OpInfo, BackendBench, and MLIR):
pytest -q

# 4. Run official PyTorch custom operator contract checks (schema, autograd, FakeTensor, AOT):
pytest tests/test_opcheck_compliance.py -v

---

License

Released under the GNU Affero General Public License v3.0 (AGPL-3.0).