eae-runtime 0.1.0__tar.gz

eae_runtime-0.1.0/.github/workflows/ci.yml

@@ -0,0 +1,80 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+  workflow_dispatch:
+
+jobs:
+  test:
+    name: test (py${{ matrix.python-version }}, ${{ matrix.os }})
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, macos-latest, windows-latest]
+        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+        exclude:
+          # CPU-only torch wheels aren't always available same-day for the
+          # newest Python on every OS; drop flaky combos here if CI is red.
+          - os: windows-latest
+            python-version: "3.9"
+
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0  # setuptools_scm needs full history/tags to compute version
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: pip
+
+      - name: Install package + dev deps
+        run: python -m pip install -e ".[dev]"
+
+      - name: Lint (ruff)
+        run: ruff check .
+
+      - name: Type check (mypy)
+        run: mypy eae_runtime
+
+      - name: Test with coverage
+        run: pytest --cov --cov-report=xml --cov-report=term-missing
+
+      - name: Upload coverage
+        if: matrix.os == 'ubuntu-latest' && matrix.python-version == '3.12'
+        uses: actions/upload-artifact@v4
+        with:
+          name: coverage-report
+          path: coverage.xml
+
+  build-check:
+    name: build sdist + wheel
+    runs-on: ubuntu-latest
+    needs: test
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Install build tooling
+        run: python -m pip install build twine
+
+      - name: Build distributions
+        run: python -m build
+
+      - name: Validate metadata
+        run: twine check dist/*
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: dist
+          path: dist/

eae_runtime-0.1.0/.github/workflows/publish.yml

@@ -0,0 +1,75 @@
+name: Publish to PyPI
+
+on:
+  push:
+    tags:
+      - "v*.*.*"       # e.g. v0.1.0 -> triggers a real release
+  release:
+    types: [published] # also allow triggering from a GitHub Release
+  workflow_dispatch:    # manual "just in case" button, publishes to TestPyPI only
+
+jobs:
+  build:
+    name: Build distribution
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0  # setuptools_scm needs tags to compute the version
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Install build tooling
+        run: python -m pip install build twine
+
+      - name: Build sdist + wheel
+        run: python -m build
+
+      - name: Check metadata
+        run: twine check dist/*
+
+      - name: Show computed version
+        run: python -c "import importlib.metadata as m; print(m.version('eae-runtime'))" || true
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+  publish-testpypi:
+    name: Publish to TestPyPI (manual trigger only)
+    if: github.event_name == 'workflow_dispatch'
+    needs: build
+    runs-on: ubuntu-latest
+    environment:
+      name: testpypi
+      url: https://test.pypi.org/p/eae-runtime
+    permissions:
+      id-token: write  # required for OIDC trusted publishing
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+      - uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          repository-url: https://test.pypi.org/legacy/
+
+  publish-pypi:
+    name: Publish to PyPI
+    if: startsWith(github.ref, 'refs/tags/v') || github.event_name == 'release'
+    needs: build
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/eae-runtime
+    permissions:
+      id-token: write  # required for OIDC trusted publishing
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+      - uses: pypa/gh-action-pypi-publish@release/v1

eae_runtime-0.1.0/.gitignore

@@ -0,0 +1,28 @@
+# Local virtualenvs
+eae_env/
+.venv/
+venv/
+
+# Python build artifacts
+__pycache__/
+*.pyc
+*.pyo
+*.egg-info/
+build/
+dist/
+*.egg
+
+# setuptools_scm-generated version file (derived from git tags at build time)
+eae_runtime/_version.py
+
+# Test / coverage / type-checker caches
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+.coverage
+htmlcov/
+
+# Editors / OS
+.vscode/
+.idea/
+.DS_Store

eae_runtime-0.1.0/LICENSE

@@ -0,0 +1,13 @@
+Copyright 2026 Vladimer Khasia
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.

eae_runtime-0.1.0/PKG-INFO

@@ -0,0 +1,229 @@
+Metadata-Version: 2.4
+Name: eae-runtime
+Version: 0.1.0
+Summary: A programmable, block-local reverse-mode execution engine that orchestrates PyTorch's local autograd (Explicit Adjoint Exposure runtime).
+Author: Vladimer Khasia
+License-Expression: Apache-2.0
+Project-URL: Homepage, https://github.com/VladimerKhasia/eae_runtime
+Project-URL: Repository, https://github.com/VladimerKhasia/eae_runtime
+Project-URL: Issues, https://github.com/VladimerKhasia/eae_runtime/issues
+Project-URL: Changelog, https://github.com/VladimerKhasia/eae_runtime/blob/main/CHANGELOG.md
+Keywords: pytorch,autograd,backpropagation,transformer,deep-learning,training-runtime,gradient-compression,distributed-training
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Typing :: Typed
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.0; extra == "dev"
+Requires-Dist: ruff>=0.6; extra == "dev"
+Requires-Dist: mypy>=1.10; extra == "dev"
+Requires-Dist: build>=1.0; extra == "dev"
+Requires-Dist: twine>=5.0; extra == "dev"
+Dynamic: license-file
+
+# A programmable reverse-mode runtime for PyTorch
+
+A production-quality runtime that replaces PyTorch's monolithic
+`backward()` with a programmable, block-local reverse-mode execution engine
+implementing **Explicit Adjoint Exposure (EAE)**. It builds on the the top of [EAE](https://github.com/VladimerKhasia/eae) and turnes reverse-mode execution itself into schedulable runtime. 
+
+The runtime is **not** a new autodiff engine. It orchestrates many local
+autograd computations using PyTorch's existing local autograd:
+
+* **PyTorch owns:** local autograd, VJP computation, kernels.
+* **The runtime owns:** forward scheduling, reverse scheduling,
+  reconstruction, adjoint lifecycle, memory lifecycle, pass execution,
+  backend selection, communication scheduling.
+
+Users write ordinary `nn.Sequential` / `nn.Module` models. No custom tensor
+type, no compiler, no FX, no PyTorch internals modified.
+
+## Install
+
+```bash
+pip install -e .
+```
+
+## Quick start
+
+```python
+import torch
+import torch.nn as nn
+from eae_runtime import EAERuntime, RuntimeConfig
+from eae_runtime.passes import ClipPass, Int8QuantizationPass
+
+model = nn.Sequential(
+    nn.Linear(64, 64), nn.ReLU(),
+    nn.Linear(64, 64), nn.ReLU(),
+    nn.Linear(64, 10),
+)
+optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+
+config = RuntimeConfig(
+    scheduler="sequential",          # "sequential" | "async" | "pipeline" | "distributed"
+    memory="pool",                   # "pool" | "none"
+    backend="auto",                  # "auto" | "cpu" | "cuda" | "triton" | "rocm"
+    passes=[Int8QuantizationPass(), ClipPass(max_norm=1.0)],
+    boundary_offload=False,          # move x0..xL to CPU between fwd/bwd
+    boundary_precision=None,         # e.g. torch.float16 to shrink boundary storage
+)
+
+runtime = EAERuntime(model, optimizer, config)
+
+x = torch.randn(32, 64)
+target = torch.randint(0, 10, (32,))
+loss = runtime.train_step(x, lambda out: nn.functional.cross_entropy(out, target))
+```
+
+## Architecture
+
+```
+User Model
+   |
+Block Decomposer            (blocks.py)
+   |
+Forward Executor            (forward_executor.py)   -> detached forward, no autograd graph
+   |
+Boundary State Store        (boundary_store.py)     -> stores only x0..xL
+   |
+Reverse Scheduler           (schedulers/)           -> pluggable: sequential/async/pipeline/distributed
+   |            \
+   |             \
+Reconstruction    Memory Manager                     (reconstruction.py / memory.py)
+Engine                 |            
+   |                   |             
+Adjoint Pipeline   Backend Manager                   (pipeline.py / backend.py)
+   |
+Local Autograd (PyTorch)
+   |
+CPU / CUDA
+```
+
+### The Adjoint Pipeline (the one addition beyond the base spec)
+
+Every reverse step exchanges an `AdjointState` — never a raw tensor —
+through a programmable pipeline:
+
+```
+AdjointState -> Pass1 -> Pass2 -> ... -> PassN -> next local VJP
+```
+
+A researcher working on gradient compression, synthetic gradients,
+error-feedback, distributed communication, or adaptive optimization only
+implements a new `EAEPass`. They never touch the reconstruction engine,
+scheduler, or memory manager.
+
+```python
+from eae_runtime.passes import EAEPass
+
+class MyPass(EAEPass):
+    name = "MyPass"
+    def process(self, adjoint, context=None):
+        new = adjoint.clone()
+        new.tensor = new.tensor * 0.9   # e.g. decay
+        return new
+```
+
+### Extension points
+
+| Want to add...          | Subclass                | Where                     |
+|--------------------------|--------------------------|----------------------------|
+| A gradient transform     | `EAEPass`                 | `eae_runtime/passes/`     |
+| A reverse scheduling policy | `BaseScheduler`        | `eae_runtime/schedulers/` |
+| A memory allocation strategy | `MemoryPolicy`        | `eae_runtime/memory.py`   |
+| A new backend             | extend `BackendManager`  | `eae_runtime/backend.py`  |
+
+None of these require modifying the runtime core.
+
+### Built-in passes
+
+`ClipPass`, `FP16Pass`, `Int8QuantizationPass` (aliased as `FP8Pass`),
+`SyntheticGradientPass`, `RegularizationPass`, `GaussianNoisePass`,
+`LoggingPass`.
+
+### Built-in schedulers
+
+* `SequentialScheduler` — the reference implementation; strict in-order
+  reconstruct → VJP → pipeline → free.
+* `AsyncScheduler` — overlaps boundary-activation prefetch (useful with
+  `boundary_offload=True`) with reconstruction compute; numerically
+  identical to `SequentialScheduler`.
+* `PipelineScheduler` — splits the batch into microbatches and runs a full
+  reverse pass per microbatch, accumulating gradients (single-process
+  emulation of pipeline-parallel microbatching).
+* `DistributedScheduler` — assigns block ranges to `torch.distributed`
+  ranks and communicates adjoints across rank boundaries; falls back to
+  `SequentialScheduler` (with a warning) when running single-process.
+
+### Events
+
+Every component emits structured events (`ForwardStarted`,
+`BlockReconstructed`, `AdjointCreated`, `MemoryAllocated`,
+`SchedulerStep`, `PassApplied`, ...) on an `EventBus`. Subscribe without
+touching runtime code:
+
+```python
+runtime.event_bus.subscribe("BlockReconstructed", lambda e: print(e.payload))
+```
+
+Set `RuntimeConfig(log_level="DEBUG")` to also pipe every event through
+Python's structured JSON logger.
+
+### Profiling
+
+```python
+loss = runtime.train_step(x, loss_fn)
+print(runtime.profiler.report())
+# {'forward': {...}, 'reconstruct:Linear_0': {...}, 'pass:ClipPass': {...}, ...}
+```
+
+## Testing
+
+```bash
+pip install -e ".[dev]"
+pytest
+```
+
+The suite (`tests/`) covers:
+
+* `AdjointState` API (norm, statistics, quantize/dequantize, compress, clone, detach)
+* Event bus pub/sub, wildcard subscribers, recording
+* Memory manager: pool reuse, null policy, allocator correctness, leak checks
+* Boundary store: precision downcast, CPU offload, ownership semantics
+* Reconstruction engine vs. manual VJP (`torch.autograd.grad`) reference
+* **Gradient equivalence**: every scheduler vs. plain `model.backward()`,
+  across float32, boundary-offloaded, and microbatched configurations
+* All built-in passes, including a full multi-pass pipeline
+* All built-in schedulers, plus a user-defined custom scheduler
+* Backend manager device/dtype resolution and graceful CUDA→CPU fallback
+* Profiler timing aggregation
+* Mixed precision (compute dtype + boundary storage dtype) numerical bounds
+* Memory-leak / allocator-stability checks over many training steps
+* A full Transformer-style block stack (`nn.MultiheadAttention` + FFN,
+  pre-norm, residual) trained end-to-end with a full pass pipeline
+* Real two-process `torch.distributed` (gloo) correctness check for
+  `DistributedScheduler` (skips gracefully if the sandbox blocks socket use)
+
+## Non-goals by design
+
+* Automatic partitioning of arbitrary `nn.Module` graphs — blocks are
+  specified manually or via `BlockDecomposer.from_sequential`/`from_list`.
+* Replacing or modifying PyTorch's autograd internals.
+* A new differentiation algorithm — local VJP is always PyTorch's own.
+* A new compiler or graph IR.
+* Day-one support for every architecture — the runtime targets
+  Transformer-style models with repeated block structure first.

eae_runtime-0.1.0/README.md

@@ -0,0 +1,192 @@
+# A programmable reverse-mode runtime for PyTorch
+
+A production-quality runtime that replaces PyTorch's monolithic
+`backward()` with a programmable, block-local reverse-mode execution engine
+implementing **Explicit Adjoint Exposure (EAE)**. It builds on the the top of [EAE](https://github.com/VladimerKhasia/eae) and turnes reverse-mode execution itself into schedulable runtime. 
+
+The runtime is **not** a new autodiff engine. It orchestrates many local
+autograd computations using PyTorch's existing local autograd:
+
+* **PyTorch owns:** local autograd, VJP computation, kernels.
+* **The runtime owns:** forward scheduling, reverse scheduling,
+  reconstruction, adjoint lifecycle, memory lifecycle, pass execution,
+  backend selection, communication scheduling.
+
+Users write ordinary `nn.Sequential` / `nn.Module` models. No custom tensor
+type, no compiler, no FX, no PyTorch internals modified.
+
+## Install
+
+```bash
+pip install -e .
+```
+
+## Quick start
+
+```python
+import torch
+import torch.nn as nn
+from eae_runtime import EAERuntime, RuntimeConfig
+from eae_runtime.passes import ClipPass, Int8QuantizationPass
+
+model = nn.Sequential(
+    nn.Linear(64, 64), nn.ReLU(),
+    nn.Linear(64, 64), nn.ReLU(),
+    nn.Linear(64, 10),
+)
+optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+
+config = RuntimeConfig(
+    scheduler="sequential",          # "sequential" | "async" | "pipeline" | "distributed"
+    memory="pool",                   # "pool" | "none"
+    backend="auto",                  # "auto" | "cpu" | "cuda" | "triton" | "rocm"
+    passes=[Int8QuantizationPass(), ClipPass(max_norm=1.0)],
+    boundary_offload=False,          # move x0..xL to CPU between fwd/bwd
+    boundary_precision=None,         # e.g. torch.float16 to shrink boundary storage
+)
+
+runtime = EAERuntime(model, optimizer, config)
+
+x = torch.randn(32, 64)
+target = torch.randint(0, 10, (32,))
+loss = runtime.train_step(x, lambda out: nn.functional.cross_entropy(out, target))
+```
+
+## Architecture
+
+```
+User Model
+   |
+Block Decomposer            (blocks.py)
+   |
+Forward Executor            (forward_executor.py)   -> detached forward, no autograd graph
+   |
+Boundary State Store        (boundary_store.py)     -> stores only x0..xL
+   |
+Reverse Scheduler           (schedulers/)           -> pluggable: sequential/async/pipeline/distributed
+   |            \
+   |             \
+Reconstruction    Memory Manager                     (reconstruction.py / memory.py)
+Engine                 |            
+   |                   |             
+Adjoint Pipeline   Backend Manager                   (pipeline.py / backend.py)
+   |
+Local Autograd (PyTorch)
+   |
+CPU / CUDA
+```
+
+### The Adjoint Pipeline (the one addition beyond the base spec)
+
+Every reverse step exchanges an `AdjointState` — never a raw tensor —
+through a programmable pipeline:
+
+```
+AdjointState -> Pass1 -> Pass2 -> ... -> PassN -> next local VJP
+```
+
+A researcher working on gradient compression, synthetic gradients,
+error-feedback, distributed communication, or adaptive optimization only
+implements a new `EAEPass`. They never touch the reconstruction engine,
+scheduler, or memory manager.
+
+```python
+from eae_runtime.passes import EAEPass
+
+class MyPass(EAEPass):
+    name = "MyPass"
+    def process(self, adjoint, context=None):
+        new = adjoint.clone()
+        new.tensor = new.tensor * 0.9   # e.g. decay
+        return new
+```
+
+### Extension points
+
+| Want to add...          | Subclass                | Where                     |
+|--------------------------|--------------------------|----------------------------|
+| A gradient transform     | `EAEPass`                 | `eae_runtime/passes/`     |
+| A reverse scheduling policy | `BaseScheduler`        | `eae_runtime/schedulers/` |
+| A memory allocation strategy | `MemoryPolicy`        | `eae_runtime/memory.py`   |
+| A new backend             | extend `BackendManager`  | `eae_runtime/backend.py`  |
+
+None of these require modifying the runtime core.
+
+### Built-in passes
+
+`ClipPass`, `FP16Pass`, `Int8QuantizationPass` (aliased as `FP8Pass`),
+`SyntheticGradientPass`, `RegularizationPass`, `GaussianNoisePass`,
+`LoggingPass`.
+
+### Built-in schedulers
+
+* `SequentialScheduler` — the reference implementation; strict in-order
+  reconstruct → VJP → pipeline → free.
+* `AsyncScheduler` — overlaps boundary-activation prefetch (useful with
+  `boundary_offload=True`) with reconstruction compute; numerically
+  identical to `SequentialScheduler`.
+* `PipelineScheduler` — splits the batch into microbatches and runs a full
+  reverse pass per microbatch, accumulating gradients (single-process
+  emulation of pipeline-parallel microbatching).
+* `DistributedScheduler` — assigns block ranges to `torch.distributed`
+  ranks and communicates adjoints across rank boundaries; falls back to
+  `SequentialScheduler` (with a warning) when running single-process.
+
+### Events
+
+Every component emits structured events (`ForwardStarted`,
+`BlockReconstructed`, `AdjointCreated`, `MemoryAllocated`,
+`SchedulerStep`, `PassApplied`, ...) on an `EventBus`. Subscribe without
+touching runtime code:
+
+```python
+runtime.event_bus.subscribe("BlockReconstructed", lambda e: print(e.payload))
+```
+
+Set `RuntimeConfig(log_level="DEBUG")` to also pipe every event through
+Python's structured JSON logger.
+
+### Profiling
+
+```python
+loss = runtime.train_step(x, loss_fn)
+print(runtime.profiler.report())
+# {'forward': {...}, 'reconstruct:Linear_0': {...}, 'pass:ClipPass': {...}, ...}
+```
+
+## Testing
+
+```bash
+pip install -e ".[dev]"
+pytest
+```
+
+The suite (`tests/`) covers:
+
+* `AdjointState` API (norm, statistics, quantize/dequantize, compress, clone, detach)
+* Event bus pub/sub, wildcard subscribers, recording
+* Memory manager: pool reuse, null policy, allocator correctness, leak checks
+* Boundary store: precision downcast, CPU offload, ownership semantics
+* Reconstruction engine vs. manual VJP (`torch.autograd.grad`) reference
+* **Gradient equivalence**: every scheduler vs. plain `model.backward()`,
+  across float32, boundary-offloaded, and microbatched configurations
+* All built-in passes, including a full multi-pass pipeline
+* All built-in schedulers, plus a user-defined custom scheduler
+* Backend manager device/dtype resolution and graceful CUDA→CPU fallback
+* Profiler timing aggregation
+* Mixed precision (compute dtype + boundary storage dtype) numerical bounds
+* Memory-leak / allocator-stability checks over many training steps
+* A full Transformer-style block stack (`nn.MultiheadAttention` + FFN,
+  pre-norm, residual) trained end-to-end with a full pass pipeline
+* Real two-process `torch.distributed` (gloo) correctness check for
+  `DistributedScheduler` (skips gracefully if the sandbox blocks socket use)
+
+## Non-goals by design
+
+* Automatic partitioning of arbitrary `nn.Module` graphs — blocks are
+  specified manually or via `BlockDecomposer.from_sequential`/`from_list`.
+* Replacing or modifying PyTorch's autograd internals.
+* A new differentiation algorithm — local VJP is always PyTorch's own.
+* A new compiler or graph IR.
+* Day-one support for every architecture — the runtime targets
+  Transformer-style models with repeated block structure first.

eae_runtime-0.1.0/eae_runtime/__init__.py

@@ -0,0 +1,80 @@
+"""
+EAE Runtime: a standalone, programmable, block-local reverse-mode execution
+engine that sits *above* PyTorch and orchestrates many local autograd
+computations, instead of replacing PyTorch's autodiff engine.
+
+Quick start:
+
+    import torch.nn as nn
+    from eae_runtime import EAERuntime, RuntimeConfig
+    from eae_runtime.passes import ClipPass, GaussianNoisePass
+
+    model = nn.Sequential(
+        nn.Linear(64, 64), nn.ReLU(),
+        nn.Linear(64, 64), nn.ReLU(),
+        nn.Linear(64, 10),
+    )
+    optimizer = torch.optim.SGD(model.parameters(), lr=1e-2)
+    config = RuntimeConfig(
+        scheduler="sequential",
+        memory="pool",
+        backend="auto",
+        passes=[ClipPass(max_norm=1.0)],
+    )
+    runtime = EAERuntime(model, optimizer, config)
+    loss = runtime.train_step(x, lambda out: nn.functional.mse_loss(out, target))
+"""
+
+from .adjoint import AdjointState
+from .backend import BackendManager
+from .blocks import BlockDecomposer, EAEBlock
+from .boundary_store import BoundaryStore
+from .config import RuntimeConfig
+from .events import Event, EventBus, EventType
+from .forward_executor import ForwardExecutor
+from .memory import MemoryManager, MemoryPolicy, NullMemoryPolicy, PoolMemoryPolicy
+from .pipeline import AdjointPipeline, PassManager
+from .profiler import Profiler
+from .reconstruction import ReconstructionEngine
+from .runtime import EAERuntime
+from .schedulers import (
+    AsyncScheduler,
+    BaseScheduler,
+    DistributedScheduler,
+    PipelineScheduler,
+    SequentialScheduler,
+    build_scheduler,
+)
+
+try:
+    from ._version import version as __version__
+except ImportError:  # pragma: no cover - only hit in a non-built source tree
+    __version__ = "0.0.0+unknown"
+
+__all__ = [
+    "AdjointState",
+    "BackendManager",
+    "BlockDecomposer",
+    "EAEBlock",
+    "BoundaryStore",
+    "RuntimeConfig",
+    "Event",
+    "EventBus",
+    "EventType",
+    "ForwardExecutor",
+    "MemoryManager",
+    "MemoryPolicy",
+    "NullMemoryPolicy",
+    "PoolMemoryPolicy",
+    "AdjointPipeline",
+    "PassManager",
+    "Profiler",
+    "ReconstructionEngine",
+    "EAERuntime",
+    "AsyncScheduler",
+    "BaseScheduler",
+    "DistributedScheduler",
+    "PipelineScheduler",
+    "SequentialScheduler",
+    "build_scheduler",
+]