belfort-ml 2026.10.8.dev0__tar.gz

belfort_ml-2026.10.8.dev0/.gitignore

@@ -0,0 +1,11 @@
+__pycache__/
+*.py[cod]
+*.egg-info/
+.mypy_cache/
+.pytest_cache/
+.ruff_cache/
+.coverage
+htmlcov/
+.venv/
+build/
+dist/

belfort_ml-2026.10.8.dev0/PKG-INFO

@@ -0,0 +1,9 @@
+Metadata-Version: 2.4
+Name: belfort-ml
+Version: 2026.10.8.dev0
+Summary: Belfort ML client for the Belfort CaaS gateway
+Requires-Python: <3.13,>=3.12
+Requires-Dist: belfort-ml-bucket==2026.10.8.dev0
+Requires-Dist: belfort-torch-mlir==2026.6.25.dev0
+Requires-Dist: httpx==0.28.1
+Requires-Dist: torch==2.13.0

belfort_ml-2026.10.8.dev0/README.md

@@ -0,0 +1,228 @@
+# Belfort ML Python Client
+
+`belfort_ml`: an httpx client for the Belfort CaaS gateway. A model provider
+uses it to compile a PyTorch model, register it, and run inference on it while
+it stays encrypted.
+
+## Install
+
+Not published. It is a `uv` workspace member, so run `uv sync` from the repo
+root.
+
+## Configuration
+
+```python
+import belfort_ml
+
+client = belfort_ml.MP_Client()  # reads the environment
+client = belfort_ml.MP_Client(api_key="blf_live_...")  # or pass it in
+```
+
+| Argument | Environment | Default |
+|---|---|---|
+| `api_key` | `BELFORT_API_KEY` | required |
+| `base_url` | `BELFORT_BASE_URL` | `https://eml.belfortlabs.com/api/v1` |
+| `timeouts` | — | `Timeouts()` |
+| `max_retries` | — | `3` |
+
+Explicit arguments beat the environment. A missing key raises
+`ConfigurationError` at construction. `Timeouts` is per operation class:
+`standard` 30s, `upload_part` 300s, `inference` 120s, `poll` 10s.
+
+## Quick start
+
+```python
+with belfort_ml.MP_Client() as client:
+    account = client.account.get()
+    print(account["mp_id"], account["credits"]["available"])
+```
+
+The client holds a connection pool for its lifetime, so use it as a context
+manager or call `close()`.
+
+`client.account` and `client.credits` return plain dicts. `client.models` and
+`client.deploys` return typed objects — `Model`, `ClientKeySet`, `Deploy`,
+`LoadedKeys`, `InferenceResult` — that carry their own methods, so
+`model.deploy()` works without passing ids around.
+
+## The gateway surface
+
+| Method | Endpoint |
+|---|---|
+| `client.account.get()` | `GET /account` |
+| `client.credits.get()` | `GET /credits` |
+| `client.credits.transactions(cursor=, limit=, kind=)` | `GET /credits/transactions` |
+| `client.models.create(artifacts, name=)` | `POST /models`, S3 PUTs, `POST /models/{m_id}/compile`, then polls |
+| `client.models.get(id)` / `.list()` | `GET /models/{m_id}` / `GET /models` |
+| `model.refresh()` / `.delete()` | `GET .../status` / `DELETE /models/{m_id}` |
+| `model.fetch_artifacts(dir, bucket_url=)` | Reads the compiled artifacts from the bucket |
+| `model.create_client(label, ttl_seconds=)` | `POST /models/{m_id}/client-tokens`: the token one end client runs on |
+| `model.artifact_downloads()` | One presigned GET per artifact of the model's compile |
+| `model.clients(status=)` / `model.delete_client(c_id)` | `GET /models/{m_id}/clients`, `DELETE /models/{m_id}/clients/{c_id}` |
+| `model.deploy(key_set=, max_lifetime=)` | `POST /deploys` for the key set it serves, then polls |
+| `client.deploys.list()` / `.get(id)` / `.stop_all()` | `GET /deploys`, `GET /deploys/{d_id}`, `DELETE` each running one |
+| `deploy.refresh()` / `.extend(n)` / `.stop()` | `GET .../status` / `POST .../extend` / `DELETE /deploys/{d_id}` |
+| `deploy.load_keys()` / `.unload_keys(k)` | `POST .../keys` for the deploy's key set / `DELETE .../keys/{k_id}` |
+| `deploy.infer(loaded, bytes)` / `.infer_batch(...)` | `POST /deploys/{d_id}/infer` |
+
+Nothing calls `GET /jobs/{job_id}`. Every `wait=True` path polls the resource's
+own `/status` route: 2s backing off to 10s, then `belfort_ml.TimeoutError`. A ^C
+while a deploy starts prints the deploy id and that the instance is still
+billing, then re-raises.
+
+`deploy(...)` is a context manager. Exit stops the instance whether the block
+completed or raised, and a deploy someone else already stopped is not an error.
+
+`client.credits.transactions` returns one page, newest first, keyset-paginated
+on `tx_id`; follow `next_cursor` yourself. The models and deploys lists follow
+cursors lazily.
+
+Uploads go straight to S3 through presigned multipart URLs, without the Belfort
+credential. The file is split into the presigned part count, with 4 parts in
+flight, each retried on its own and streamed from disk.
+
+## Compiling locally
+
+```python
+from belfort_ml._compile._export import load_calibration_data, load_pytorch_model
+
+model = load_pytorch_model("smallnet.pt")  # eval mode, CPU
+calibration = load_calibration_data(  # 64 examples by default
+    "calibration.npz", key="inputs", input_shape=(1, 3, 32, 32)
+)
+artifacts = belfort_ml.compile(
+    model, calibration[0], "smallnet", "./build", calibration_inputs=calibration
+)
+```
+
+`compile` writes `build/smallnet_torch.mlir` and returns the `Artifacts` that
+`client.models.create` takes: the file and the parameters it declares. No
+network is involved.
+
+`name` becomes the generated C++ code's name, so it must be a C++ identifier:
+ASCII letters, digits and underscores, not starting with a digit, at most 200
+characters, with no `__` and no leading `_` before an uppercase letter. It may
+not be a C++ keyword, `std`, `heir`, `detail` or one of the native build's
+macros (`PERSEUS_MODEL`, `PERSEUS_HEADER`, `PERSEUS_SERVER`, `DEFAULT_STREAM`)
+or a standard-library macro (`NULL`, `EOF`, `errno`, `stdin`, `stdout`, `stderr`). Any other name raises `ValueError`
+before export.
+
+`example_input` fixes the shapes. The graph is specialized to it, so a deployed
+model serves that input shape only. A model the frontend cannot lower raises
+`CompileError` carrying the frontend's message.
+
+Two annotations travel with the IR. Every input carries `{secret.secret}` unless
+the model marks a subset with `belfort_ml.Secret`, as
+[`examples/mlp/mlp.py`](../examples/mlp/mlp.py) does; and each op the compiler
+must approximate carries the value range the calibration examples put through
+it. Only those bounds reach the IR, never the examples.
+
+The end client is its own component, [`belfort-ml-client/`](../belfort-ml-client).
+
+## Reading the compiled artifacts
+
+Perseus writes under `{mp_id}/{m_id}/compiles/{compile_id}/`. The gateway
+reports a status and the chosen parameters; the files themselves are read from
+the bucket:
+
+```python
+compiled = model.fetch_artifacts(
+    "./artifacts", bucket_url="http://localhost:9000/belfort-caas-models"
+)
+compiled.compile_id  # the model's compile
+compiled.params  # what the compiler actually chose
+compiled.entry_function  # the exported function the C++ is named after
+compiled.files  # every downloaded file, tree preserved
+```
+
+| Argument | Environment | Default |
+|---|---|---|
+| `bucket_url` | `BELFORT_BUCKET_URL` | required |
+| `access_key_id` | `BELFORT_BUCKET_ACCESS_KEY_ID` | boto3's own chain |
+| `secret_access_key` | `BELFORT_BUCKET_SECRET_ACCESS_KEY` | boto3's own chain |
+| `region` | `BELFORT_BUCKET_REGION` | `us-east-1` |
+
+`bucket_url` names the endpoint and the bucket together, as
+[`bucket/`](../bucket) describes; a URL carrying a key prefix raises
+`ConfigurationError`.
+
+It reads the model's compile, `model.active_compile_id`. Pass `store=` a
+`belfort_bucket.Bucket` to reuse one reader across several models.
+
+## The end client
+
+`model.create_client(label)` opens a key set for one end client and returns a
+`ClientToken`. Hand its `token` to the client on your own channel; the client
+then talks to the gateway itself, with no bucket credential and no API key:
+
+```bash
+belfort-ml-client infer --token "$BELFORT_CLIENT_TOKEN" --input input.json
+```
+
+The client downloads and links its code, generates its key pair, uploads the
+evaluation keys, and waits for the gateway to load them onto the deploy started
+for its key set with `model.deploy(key_set=minted, max_lifetime=...)`. [`belfort-ml-client`](../belfort-ml-client) describes that side, and
+[`e2e/`](../tests/e2e) runs both.
+
+A key set travels as a directory, one object per evaluation key: `<n>.evk` and
+one `<n>.s2d`.
+
+## Errors
+
+Everything descends from `BelfortMLError`, and local failures separate from
+gateway failures.
+
+| Raised locally | When |
+|---|---|
+| `ConfigurationError` | No API key, or a `base_url` that is not http(s) |
+| `CompileError` | Local or gateway-side compile failure |
+| `UploadError` | An S3 part failed after retries |
+| `DownloadError` | Reading the bucket failed, or it holds no such compile |
+| `KeyMaterialError` | A key directory holds something that may not be published |
+| `TimeoutError` | Poll deadline exceeded |
+| `ConnectionError` | Gateway unreachable after retries; httpx's error rides along as `__cause__` |
+
+`APIError` covers anything the gateway returned, carrying `message`, `code`,
+`detail`, `status_code` and `request_id`. Quote `request_id` when you report a
+failure; it locates the request in the gateway logs.
+
+| Status | Class | Extra |
+|---|---|---|
+| 400, 413 | `ValidationError` | |
+| 401 | `AuthenticationError` | |
+| 402 | `InsufficientCredits` | `.required`, `.available` |
+| 403 | `PermissionError` | |
+| 404 | `NotFoundError` | |
+| 409 | `ConflictError` | `DeployLimitReached` for `deploy_limit_reached`, with `.limit`, `.running` |
+| 429 | `RateLimitError` | `.retry_after` |
+| 503 | `ServiceUnavailable` | |
+
+A response that is not the gateway's `{"error": {...}}` envelope — a proxy page,
+a crash — degrades to the raw text.
+
+## Retries, idempotency and logging
+
+Connection errors, timeouts, 429 and 500/502/504 are retried `max_retries`
+times, with exponential backoff and full jitter capped at 30s. A 429's
+`Retry-After` wins over the computed backoff. 4xx is never retried, and neither
+is inference. A read or write timeout can come after the gateway acted on the
+request, so after one only GET and HEAD are sent again.
+
+Idempotent requests carry one `Idempotency-Key` per logical operation, reused
+across that operation's retries. The gateway does not honour it yet.
+
+Logging goes to the standard `belfort_ml` logger and is silent without a
+handler. DEBUG lines carry method, path, status and request id only, never
+headers or bodies.
+
+## Examples and tests
+
+[`examples/`](./examples/) covers the account, credits and error handling.
+
+```bash
+cd gateway && docker compose up -d minio    # the artifact tests need it
+uv run pytest                                # from belfort-ml/
+```
+
+CI runs this package from
+[`.github/workflows/test-packages.yml`](../.github/workflows/test-packages.yml).

belfort_ml-2026.10.8.dev0/belfort_ml/__init__.py

@@ -0,0 +1,72 @@
+__version__ = "2026.10.8.dev0"
+
+from belfort_bucket import CompiledArtifacts
+
+from belfort_ml._compile import compile  # noqa: A004
+from belfort_ml._compile._annotations import Secret
+from belfort_ml.errors import (
+    APIError,
+    AuthenticationError,
+    BelfortMLError,
+    CompileError,
+    ConfigurationError,
+    ConflictError,
+    ConnectionError,
+    DeployLimitReached,
+    DownloadError,
+    InsufficientCredits,
+    KeyMaterialError,
+    NotFoundError,
+    PermissionError,
+    RateLimitError,
+    ServiceUnavailable,
+    TimeoutError,
+    UploadError,
+    ValidationError,
+)
+from belfort_ml.gateway_connector.client import MP_Client
+from belfort_ml.gateway_connector.deploys import (
+    Deploy,
+    InferenceResult,
+    LoadedKeys,
+)
+from belfort_ml.gateway_connector.models import (
+    ClientKeySet,
+    ClientToken,
+    Model,
+)
+from belfort_ml.types import Artifacts, CompileParams
+
+__all__ = [
+    "APIError",
+    "Artifacts",
+    "AuthenticationError",
+    "MP_Client",
+    "ClientKeySet",
+    "CompileError",
+    "CompileParams",
+    "CompiledArtifacts",
+    "Deploy",
+    "DeployLimitReached",
+    "ClientToken",
+    "InferenceResult",
+    "KeyMaterialError",
+    "LoadedKeys",
+    "Model",
+    "ConfigurationError",
+    "ConflictError",
+    "ConnectionError",
+    "DownloadError",
+    "InsufficientCredits",
+    "BelfortMLError",
+    "NotFoundError",
+    "PermissionError",
+    "RateLimitError",
+    "Secret",
+    "ServiceUnavailable",
+    "TimeoutError",
+    "UploadError",
+    "ValidationError",
+    "__version__",
+    "compile",
+]

belfort_ml-2026.10.8.dev0/belfort_ml/_compile/__init__.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+import re
+from pathlib import Path
+from typing import Any
+
+from belfort_ml._compile._export import SampleInput, export_torch_mlir
+from belfort_ml.types import Artifacts
+
+# `name` names the C++ Perseus generates (`heir::generated::<name>`) and its
+# artifacts, so it must be a C++ identifier that is neither a keyword, a
+# namespace that code refers to, nor a macro the native build or the standard
+# headers define, and short enough that `<name>_cheddar.mlir` fits a 255-byte
+# file name.
+_NAME = re.compile(r"[A-Za-z_][A-Za-z0-9_]{0,199}")
+_RESERVED_NAMES = frozenset(
+    """
+    alignas alignof and and_eq asm auto bitand bitor bool break case catch char
+    char8_t char16_t char32_t class compl concept const consteval constexpr
+    constinit const_cast continue co_await co_return co_yield decltype default
+    delete do double dynamic_cast else enum explicit export extern false float
+    for friend goto if inline int long mutable namespace new noexcept not not_eq
+    nullptr operator or or_eq private protected public register reinterpret_cast
+    requires return short signed sizeof static static_assert static_cast struct
+    switch template this thread_local throw true try typedef typeid typename
+    union unsigned using virtual void volatile wchar_t while xor xor_eq
+    detail heir std
+    DEFAULT_STREAM PERSEUS_HEADER PERSEUS_MODEL PERSEUS_SERVER
+    EOF NULL errno stderr stdin stdout
+    """.split()
+)
+
+
+def _validate_name(name: str) -> None:
+    if not _NAME.fullmatch(name):
+        raise ValueError(
+            f"model name {name!r} may use only ASCII letters, digits and "
+            "underscores, may not start with a digit, and may be at most 200 "
+            "characters long"
+        )
+    # C++ reserves these for the implementation (`__FILE__`, `_Bool`).
+    if name in _RESERVED_NAMES or "__" in name or re.match(r"_[A-Z]", name):
+        raise ValueError(
+            f"model name {name!r} is reserved in the C++ the compiler generates"
+        )
+
+
+def compile(  # noqa: A001
+    model: Any,  # torch.nn.Module
+    example_input: SampleInput,
+    name: str,
+    output_dir: str | Path,
+    *,
+    calibration_inputs: list[SampleInput] | None = None,
+) -> Artifacts:
+    _validate_name(name)
+    mlir_path = export_torch_mlir(
+        model,
+        example_input,
+        name,
+        output_dir,
+        calibration_inputs=calibration_inputs,
+    )
+    return Artifacts(mlir_path=mlir_path)

belfort_ml-2026.10.8.dev0/belfort_ml/_compile/_annotations.py

@@ -0,0 +1,299 @@
+"""Secret-input markers and range calibration, applied inside `export_and_import`'s `annotate=` callback."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Iterable
+from typing import Any, get_type_hints
+
+import torch
+import torch.nn as nn
+from torch.export import ExportedProgram
+from torch.export.graph_signature import InputKind
+from torch_mlir.extras import fx_importer  # type: ignore[import-untyped]
+from torch_mlir.extras.annotate import annotate_arg  # type: ignore[import-untyped]
+
+from belfort_ml.errors import CompileError
+
+
+class Secret:
+    """PEP-593 marker: the annotated forward argument is a ciphertext input."""
+
+
+SECRET_ATTR = {"secret.secret": True}
+
+# Tail probability dropped at each end of a calibrated range; 0 gives the exact min/max.
+DEFAULT_LOWER_QUANTILE = 0.01
+
+# The calibrated half-range is multiplied by this.
+DEFAULT_DOMAIN_MARGIN = 2.0
+
+# Values sampled per recorded node.
+CAPACITY = 1 << 14
+
+
+def apply_forward_annotations(prog: ExportedProgram, model: nn.Module) -> int:
+    """Mark the forward arguments annotated with `Secret`; return how many."""
+    try:
+        hints = get_type_hints(model.forward, include_extras=True)
+    except (NameError, TypeError) as exc:
+        raise CompileError(
+            f"could not read the annotations of {type(model).__name__}.forward: "
+            f"{exc}. Make every annotation resolvable from the module that "
+            f"defines the model, so the `Secret` markers can be read."
+        ) from exc
+    marked = 0
+    for name, hint in hints.items():
+        if name != "return" and Secret in getattr(hint, "__metadata__", ()):
+            annotate_arg(prog, name, dict(SECRET_ATTR))
+            marked += 1
+    return marked
+
+
+def annotate_all_inputs_secret(prog: ExportedProgram) -> int:
+    """Mark every user input as a ciphertext input; return how many."""
+    n_inputs = sum(
+        1
+        for spec in prog.graph_signature.input_specs
+        if spec.kind == InputKind.USER_INPUT
+    )
+    for index in range(n_inputs):
+        annotate_arg(prog, index, dict(SECRET_ATTR))
+    return n_inputs
+
+
+class ValueSample:
+    """Reservoir sample (vectorized Algorithm R) of a stream of tensors, plus the exact min and max."""
+
+    def __init__(self, generator: torch.Generator) -> None:
+        self._generator = generator
+        self._reservoir = torch.empty(0, dtype=torch.float32)
+        self.count = 0
+        self._min = math.inf
+        self._max = -math.inf
+
+    def update(self, t: torch.Tensor) -> None:
+        values = t.detach().reshape(-1).to(torch.float32)
+        if values.numel() == 0:
+            return
+        self._min = min(self._min, float(values.min()))
+        self._max = max(self._max, float(values.max()))
+
+        free = CAPACITY - self._reservoir.numel()
+        if free > 0:
+            take = min(free, values.numel())
+            self._reservoir = torch.cat([self._reservoir, values[:take]])
+            self.count += take
+            values = values[take:]
+        if values.numel() > 0:
+            # Stream index i replaces a random slot with probability CAPACITY/(i+1).
+            index = torch.arange(self.count, self.count + values.numel())
+            accept = torch.rand(
+                values.numel(), generator=self._generator
+            ) < CAPACITY / (index + 1).to(torch.float32)
+            n_accept = int(accept.sum())
+            if n_accept > 0:
+                slots = torch.randint(
+                    0, CAPACITY, (n_accept,), generator=self._generator
+                )
+                self._reservoir[slots] = values[accept]
+            self.count += values.numel()
+
+    def quantile(self, q: float) -> float:
+        if q <= 0.0:
+            return self._min
+        if q >= 1.0:
+            return self._max
+        return float(torch.quantile(self._reservoir, q))
+
+
+# Aten ops that lower to an elementwise `linalg.generic` and that the compiler
+# approximates by a polynomial; the CKKS-exact add/sub/mul/neg are absent.
+_ELEMENTWISE_OP_NAMES = (
+    "relu", "relu6", "leaky_relu", "prelu", "gelu", "silu", "mish",
+    "elu", "celu", "selu", "sigmoid", "hardsigmoid", "hardswish",
+    "hardtanh", "hardshrink", "softshrink", "softplus", "threshold",
+    "tanh", "erf", "logit",
+    "exp", "exp2", "expm1", "log", "log1p", "log2", "log10",
+    "sqrt", "rsqrt", "reciprocal", "pow",
+    "sin", "cos", "tan", "asin", "acos", "atan", "atan2",
+    "sinh", "cosh", "asinh", "acosh", "atanh",
+    "div", "remainder", "fmod",
+    "abs", "sign", "floor", "ceil", "round", "trunc",
+    "maximum", "minimum", "clamp", "clamp_min", "clamp_max",
+    "where", "lerp",
+)  # fmt: skip
+
+# Ops whose domain is not all of R: widening could cross a singularity, so they
+# keep their observed range.
+_RESTRICTED_DOMAIN_OP_NAMES = (
+    "sqrt", "rsqrt", "reciprocal", "pow",
+    "log", "log1p", "log2", "log10", "logit",
+    "div", "remainder", "fmod",
+    "asin", "acos", "atanh", "acosh",
+)  # fmt: skip
+
+
+def _aten_ops(names: tuple[str, ...]) -> frozenset[Any]:
+    return frozenset(
+        packet
+        for name in names
+        if (packet := getattr(torch.ops.aten, name, None)) is not None
+    )
+
+
+ELEMENTWISE_ATEN_OPS = _aten_ops(_ELEMENTWISE_OP_NAMES)
+RESTRICTED_DOMAIN_ATEN_OPS = _aten_ops(_RESTRICTED_DOMAIN_OP_NAMES)
+
+
+def _targets(node: torch.fx.Node, ops: frozenset[Any]) -> bool:
+    return (
+        node.op == "call_function"
+        and getattr(node.target, "overloadpacket", None) in ops
+    )
+
+
+def lowers_to_elementwise(node: torch.fx.Node) -> bool:
+    return _targets(node, ELEMENTWISE_ATEN_OPS)
+
+
+class RangeRecorder(torch.fx.Interpreter):
+    """Interprets an exported program, sampling the values that feed each
+    elementwise-lowered op: user inputs and `call_function` results, never
+    parameters, buffers or constants.
+    """
+
+    def __init__(self, prog: ExportedProgram) -> None:
+        super().__init__(prog.graph_module)
+        self._prog = prog
+        # Seeded, so calibration is repeatable.
+        self._generator = torch.Generator().manual_seed(0)
+        self._user_inputs = {
+            spec.arg.name
+            for spec in prog.graph_signature.input_specs
+            if spec.kind == InputKind.USER_INPUT
+        }
+        # Parameters, buffers and constants: lifted to inputs the caller does not supply.
+        self._lifted = {**prog.state_dict, **prog.constants}
+        # Producers feeding an elementwise-lowered op.
+        self._monitored = {
+            src.name
+            for node in prog.graph.nodes
+            if lowers_to_elementwise(node)
+            for src in node.all_input_nodes
+        }
+        self.samples: dict[str, ValueSample] = {}
+
+    def run_node(self, n: torch.fx.Node) -> Any:
+        result = super().run_node(n)
+        recorded = n.op == "call_function" or (
+            n.op == "placeholder" and n.name in self._user_inputs
+        )
+        if (
+            n.name in self._monitored
+            and recorded
+            and isinstance(result, torch.Tensor)
+            and result.is_floating_point()
+        ):
+            sample = self.samples.setdefault(n.name, ValueSample(self._generator))
+            sample.update(result)
+        return result
+
+    def run_calibration(self, batch: torch.Tensor | tuple[Any, ...]) -> None:
+        """Run one calibration example, shaped like the export sample, in forward-argument order."""
+        user_inputs = iter((batch,) if isinstance(batch, torch.Tensor) else batch)
+        args = [
+            next(user_inputs)
+            if spec.kind == InputKind.USER_INPUT
+            else self._lifted[spec.target]
+            for spec in self._prog.graph_signature.input_specs
+        ]
+        with torch.no_grad():
+            self.run(*args)
+
+    def bounds(self, q: float) -> dict[str, tuple[float, float]]:
+        """Per-op input domains: for each elementwise-lowered node, the union of
+        its producers' [q, 1-q] output ranges. Nodes with no recorded producer —
+        one fed only by weights, say — are omitted.
+        """
+        if not 0 <= q < 0.5:
+            raise ValueError(f"tail probability q must satisfy 0 <= q < 0.5, got {q}")
+        out: dict[str, tuple[float, float]] = {}
+        for node in self._prog.graph.nodes:
+            if not lowers_to_elementwise(node):
+                continue
+            lo, hi = math.inf, -math.inf
+            for src in node.all_input_nodes:
+                sample = self.samples.get(src.name)
+                if sample is not None and sample.count > 0:
+                    lo = min(lo, sample.quantile(q))
+                    hi = max(hi, sample.quantile(1 - q))
+            if lo <= hi:
+                out[node.name] = (lo, hi)
+        return out
+
+
+def record_ranges(
+    prog: ExportedProgram,
+    calibration_data: Iterable[torch.Tensor | tuple[Any, ...]],
+    *,
+    q: float = DEFAULT_LOWER_QUANTILE,
+    margin: float = DEFAULT_DOMAIN_MARGIN,
+) -> dict[str, tuple[float, float]]:
+    """`{fx_node_name: (domain_lo, domain_hi)}` for `apply_ranges`.
+
+    The domain is the [q, 1-q] range of the values the op was evaluated at,
+    widened about its centre by `margin`; restricted-domain ops keep the raw range.
+    """
+    recorder = RangeRecorder(prog)
+    for batch in calibration_data:
+        recorder.run_calibration(batch)
+
+    nodes = {node.name: node for node in prog.graph.nodes}
+    out: dict[str, tuple[float, float]] = {}
+    for name, (lo, hi) in recorder.bounds(q).items():
+        if _targets(nodes[name], RESTRICTED_DOMAIN_ATEN_OPS):
+            out[name] = (lo, hi)
+        else:
+            centre, half = (lo + hi) / 2, (hi - lo) / 2 * margin
+            out[name] = (centre - half, centre + half)
+    return out
+
+
+def apply_ranges(prog: ExportedProgram, bounds: dict[str, tuple[float, float]]) -> int:
+    """Attach `{domain_lower, domain_upper}` to the named nodes; return how many."""
+    nodes = {node.name: node for node in prog.graph.nodes}
+    hits = 0
+    for name, (lo, hi) in bounds.items():
+        node = nodes.get(name)
+        if node is None:
+            continue
+        # Rounded, so the emitted MLIR is reproducible across machines.
+        node.meta.setdefault(fx_importer.MLIR_OP_ATTRS_META_KEY, {}).update(
+            {"domain_lower": round(float(lo), 4), "domain_upper": round(float(hi), 4)}
+        )
+        hits += 1
+    return hits
+
+
+def freeze_buffers(prog: ExportedProgram) -> int:
+    """Inline every unmutated buffer as a literal; return how many."""
+    signature = prog.graph_signature
+    mutated = set(signature.buffers_to_mutate.values())
+    # A `persistent=False` buffer lives in `constants`, not `state_dict`.
+    lifted = {**prog.state_dict, **prog.constants}
+    replacements = {
+        input_name: lifted[state_name]
+        for input_name, state_name in signature.inputs_to_buffers.items()
+        if state_name not in mutated
+    }
+
+    frozen = 0
+    for node in list(prog.graph.nodes):
+        if node.op != "placeholder" or node.name not in replacements:
+            continue
+        # The importer emits a substituted tensor as a `torch.vtensor.literal`.
+        node.replace_all_uses_with(replacements[node.name])
+        prog.graph.erase_node(node)
+        frozen += 1
+    return frozen