Unified Inference Engine: High-Performance JAX Framework

The Unified Inference Engine is a high-performance framework designed for simulation-based parameter inference, non-convex likelihood sampling, and sub-millisecond continuous-time feature extraction.

The architecture bridges the gap between scientific computing and low-latency Quantitative Finance, providing infrastructure capable of handling extreme dimensionality, heterogeneous hardware deployments (CPU-GPU zero-copy), and noisy, high-frequency environments.

🚀 Key Modules

1. EMRI Astrophysics (MCMC & Normalizing Flows)

• Domain: Gravitational Waveform Analysis for Extreme Mass-Ratio Inspirals.
• Core Technology:
  • Parallel Tempering HMC (PT-HMC): JAX-native implementation optimized for complex manifold sampling. Includes automated replica exchange logic.
  • Normalizing Flows (CNF): Amortized Simulation-Based Inference (SBI) utilizing RealNVP-based flows for fast prior narrowing.
  • Custom VJP Rules: Integration of C++ waveform backends with full Gradient/Jacobian-vector-product (JVP) support via vectorized finite differences.

2. Quant Finance Pipeline (Neural CDE Engine)

• Domain: Limit Order Book (LOB) Microstructure & Prediction.
• Core Technology:
  • Neural Controlled Differential Equations (CDEs): High-performance time-series integration on GPU, reaching sub-millisecond inference latencies.
  • Asymmetric Manifold Calibration: Proprietary Lipschitz-bounding logic for numerical stability in highly volatile regimes.
  • Recursive Checkpointing: Optimized backpropagation memory management.

3. High-Frequency Defense (Shield of Gaia)

• Domain: Sub-5ms flash-crash detection and spoofing suppression.
• Core Technology:
  • Zero-Copy Rust FFI: Low-latency data ingestion using the AF_XDP driver for kernel-bypass networking.
  • XLA Data Entanglement: Architectural implementation of hardware-level read barriers using JAX bitwise-fusion to eliminate memory tearing.

🛠️ Repository Structure

• core/:
  • model.py: Neural CDE and Manifold Calibration.
  • consistency_gate.py: XLA bitwise fusion and shock detection.
  • emri_infer/: PT-HMC engine, CNF surrogates, and FFI bridges.
• scripts/: Production pipeline and benchmark entry points.

Important

IP Protection Notice: This repository is a technical showcase of high-performance engineering architecture. All proprietary mathematical formulas, physical potentials, and exact numerical constants are omitted or replaced by abstract kernels to protect research and commercial IP.

📜 Credits

Developed as part of a high-performance computing research initiative targeting AI-accelerated physics and quant engineering.