Agent¶

laya.Agent loads one checkpoint and answers typed questions about a state. laya.load is a shortcut for Agent(...) for every backend except backend="onnx", which returns an ONNXAgent instead, and laya.RLAgent is an alias of Agent. ONNXAgent runs an exported ONNX model on CPU; import it from laya.onnx_agent. It is not an Agent subclass -- both are HookRegistrys, but only Agent holds a PyTorch model, so isinstance(laya.load(..., backend="onnx"), laya.Agent) is False.

Agent ¶

Agent(
    model_id_or_path: str = "convaiinnovations/laya",
    device: Optional[str] = None,
    token: Optional[str] = None,
    subfolder: Optional[str] = None,
    fast: bool = False,
    compile: bool = False,
    revision: Optional[str] = None,
    expected_sha256: Optional[Dict[str, str]] = None,
    lang_temperatures: Optional[Dict[str, Dict[str, Any]]] = None,
    hooks=None,
    on_predict_start=None,
    on_predict_end=None,
    hooks_raise: bool = True,
    hooks_concurrent: bool = True,
    hooks_timeout: Optional[float] = None,
    calibration: Optional[str] = None,
    backend: Optional[str] = None,
    compile_warmup: bool = True,
    compile_cache: bool = False,
    compile_mode: str = "default",
)

Bases: HookRegistry

System 1 decision model runtime: fast, non-autoregressive, calibrated decisions.

dtype is the autocast target, not the precision of every call. On MPS a call autocasts only at or above mps_amp_min_rows rows, so dtype can say float16 while a call runs in float32. dtype_for(rows) returns the precision of a call with rows rows.

Load a Laya checkpoint.

backend selects "eager", "auto", "compile" or "tilelang"; see laya.backends. It takes precedence over fast and compile. Omit it to keep those legacy flags. For ONNX use load(backend="onnx") instead.

revision optionally pins the Hub download to an explicit commit SHA/branch/tag; when omitted, huggingface_hub's normal default and existing offline cache are used. expected_sha256 ({path relative to the checkpoint dir: hexdigest}) verifies artifact integrity before any weight is parsed or executed; it is opt-in and applies to local directories too. A missing artifact raises FileNotFoundError and a digest mismatch raises ValueError; either error refuses the load.

fast=True swaps the encoder/head forward for the TileLang fast path (CUDA only, needs pip install laya[fast]); see Agent.accelerate.

compile=True runs the model under torch.compile and turns ModernBERT's encoder reference_compile on. torch.compile specializes per input shape and Laya sees a new one on almost every request, so those graphs usually cost more than they return; use it when the traffic is repetitive. fast=True takes precedence, because the TileLang path replaces the forward that would be compiled. Compiled agents run warmup() before returning; compile_warmup=False defers that work to requests or a manual warmup() call. Eager and fast agents are unchanged. compile_cache=True opts into a persistent Laya Inductor directory (process-wide), respecting any existing TORCHINDUCTOR_CACHE_DIR; see the compile engineering notes. compile_mode="reduce-overhead" opts into CUDA graphs. It can retain more GPU memory and records each new shape separately. CUDA outputs are copied before the next replay; compiled CUDA graph forwards are serialized. The default mode remains "default".

subfolder selects one checkpoint from a repo that bundles several, e.g. Agent("convaiinnovations/laya", subfolder="multilingual"). Only that subfolder is downloaded, so bundling does not cost every user the whole family.

calibration is an optional JSON path with temperature and temperature_by_options. It is applied after the checkpoint config, so a fitted map overrides shipped scalars without rewriting model.safetensors.

hooks / on_predict_start / on_predict_end observe or shape every prediction; see laya.hooks. hooks_raise=False warns and continues when a hook fails, hooks_concurrent=False serialises hooks that are not safe to run in parallel, and hooks_timeout bounds each hook call in seconds (None means no limit).

backend `property` ¶

backend: str

The active inference backend, including the legacy compile and fast flags.

backend_object `property` ¶

backend_object

The installed Backend object, or None for a legacy runtime.

set_backend ¶

set_backend(
    name: str = "auto", strict: bool = False, **options
) -> str

Switch backends; unavailable backends warn and use eager unless strict=True.

Options go to the backend constructor, e.g. warmup=False for compile or use_graphs=False for tilelang. Switching waits for in-flight inference.

accelerate ¶

accelerate(use_graphs: bool = True, strict: bool = False)

Replace the model forward with the TileLang fast path (fused GEMM/GEGLU/LayerNorm/RoPE kernels, sliding-window flash attention, 16-bit resident weights, CUDA graphs per shape bucket).

The fast path runs in the agent's autocast dtype at the time of the call (bf16 or fp16), so it matches the stock forward it replaces within rounding (see benchmarks/parity_fast.py). After changing agent.dtype, call deaccelerate() then accelerate() to rebuild it. Returns True if enabled. With strict=False any failure (no CUDA, tilelang missing) leaves the stock path in place.

warmup ¶

warmup(shapes=None) -> float

Run the forward on synthetic input of each shape now and return the seconds it took.

compile=True calls this at load unless compile_warmup=False. Extra shapes can still be warmed manually. fast=True builds its kernels and CUDA graphs per shape bucket on first use; calling this before serving moves that cost out of the first requests. With the stock forward it is a few ordinary forward passes. shapes is a list of (rows, tokens, markers); tokens are capped at the agent's max_len. Nothing is returned to or recorded for any caller, and hooks do not run.

deaccelerate ¶

deaccelerate()

Restore the stock forward.

dtype_for ¶

dtype_for(rows: int) -> torch.dtype

Precision that a forward pass with rows question rows runs in.

dtype is the autocast target, set once at load time. Whether a forward autocasts is decided per call: on MPS only at or above mps_amp_min_rows rows. This returns dtype when a forward with rows rows autocasts, and torch.float32 when it does not. A predict call runs one row per question.

predict_batch ¶

predict_batch(
    states: List[Union[str, dict, list]],
    questions: Dict[str, Dict[str, Any]],
    batch_size: Optional[int] = None,
    lang: Optional[str] = None,
    hooks=None,
    on_predict_start=None,
    on_predict_end=None,
    hooks_raise: Optional[bool] = None,
    hooks_timeout: Optional[float] = None,
    max_len: Optional[int] = None,
    head_max_len: Optional[int] = None,
    sort_by_length: bool = False,
    min_confidence: Optional[float] = None,
) -> List[Dict[str, Any]]

Evaluate the same questions over many states, packing them into shared forward passes.

This is the throughput path. system_one/predict handle one state per forward pass; on a GPU that leaves most of the batch dimension idle. predict_batch collates several states' question rows into one tensor, so a call that would take N sequential forward passes takes one (or ceil(len(states) / batch_size)), which is several times faster per decision on GPU.

Parameters:

states (List[Union[str, dict, list]]) –

A list of states (each a text string, JSON dict, or conversation turn list). The same questions are evaluated against every state.
questions (Dict[str, Dict[str, Any]]) –

Question definitions, exactly as accepted by system_one.
batch_size (Optional[int], default: None ) –

Optional cap on states per forward pass. None sends them all in one pass; set it to bound peak memory when batching many or long states.
hooks (HookArg, default: None ) –

Per-call hooks, appended after any installed on the Agent. See laya.hooks.
on_predict_start (PredictHookArg, default: None ) –

A per-call start hook. It may rewrite the state/questions or call ctx.skip(...) to short-circuit inference.
on_predict_end (PredictHookArg, default: None ) –

A per-call end hook. It may rewrite the results.
hooks_raise (Optional[bool], default: None ) –

Override the Agent's hooks_raise for this call.
hooks_timeout (Optional[float], default: None ) –

Override the Agent's hooks_timeout for this call.
max_len (Optional[int], default: None ) –

Override the agent config's max_len for this call. A start hook may also set ctx.max_len to shape the token budget.
head_max_len (Optional[int], default: None ) –

Override the agent config's head_max_len for this call. A start hook may also set ctx.head_max_len.
sort_by_length (bool, default: False ) –

Group similarly sized encoded states within windows of eight batches to reduce padding. Requires an explicit batch_size greater than one and smaller than the number of states; otherwise it has no effect. Results retain input order. This buffers up to eight batches of tokenized states instead of one. Changing batch shapes can slightly change floating-point predictions.

Returns:

List[Dict[str, Any]] –

A list of per-state result dicts, each identical in shape to system_one's output and aligned with states by index.

predict_long ¶

predict_long(
    state: Union[str, dict, list],
    questions: Dict[str, Dict[str, Any]],
    window: Optional[int] = None,
    stride: Optional[int] = None,
    aggregate: str = "auto",
    batch_size: Optional[int] = None,
    lang: Optional[str] = None,
    hooks=None,
    on_predict_start=None,
    on_predict_end=None,
    hooks_raise: Optional[bool] = None,
    hooks_timeout: Optional[float] = None,
) -> Dict[str, Any]

Evaluate questions over a state longer than the context window, scanning it in overlapping windows and aggregating per question.

system_one/predict truncate a state that exceeds max_len to a single window (the first, or for a conversation list the last), silently dropping the rest. predict_long tokenizes the state once, splits it into overlapping token windows, scores every window in shared forward passes (via predict_batch), and combines the per-window answers:

noul -> P(true) is the max over windows (the statement holds if any window supports it)
choice-> the answer from the single most-confident window, so a localized signal isn't out-voted by the many neutral windows a long document is mostly made of (averaging drowns it -- the neutral majority dominates)
score -> the level from the most-confident window, likewise

The returned probability/confidence is the deciding window's, not a calibrated number for the whole document: a noul max over many windows drifts up with the window count even with no signal, and choice can land on a confidently-neutral window when nothing in the document is decisive. Each answer therefore carries answer["window"] — the deciding window's index, token_start/token_end into the tokenized state, and the window count — so a caller can inspect the span the answer came from rather than trust the raw number. That span is the one the model read, not merely the one asked for: the window is capped at the room the questions leave, so what is handed to predict_batch is not cut short again.

A state the questions leave room for is passed straight to system_one (identical output), since system_one reads it whole; with an explicit window, a state that fits that window.

The hooks wrap the inference that answers the state, which for a document needing several windows is the one shared predict_batch over them: on_predict_start fires once, and ctx.states holds the decoded window texts in scan order -- not the caller's state, which was tokenized to produce them. Three outcomes follow from what the chain leaves behind:

ctx.skip([result]) answers the document: the payload comes back with no window attribution and usage["windows"] at 0, because nothing was scored
a scan left as this method built it: every window is scored, each answer carries answer["window"], and usage["windows"] is the window count
a rewritten scan (ctx.states replaced, in any way): the answers are aggregated over the states that were scored, but no answer carries answer["window"] -- the offsets above describe this method's windows, not the text the model read

Parameters:

window (Optional[int], default: None ) –

state tokens per window. Defaults to the checkpoint's state budget, max(64, max_len - head_max_len - 8) -- the 64 is a floor, so widening head_max_len stops shrinking the default once the budget reaches it -- and either way is capped at the room the questions leave for the state inside max_len -- the smallest room of them, because one list of windows is scored for every question. A wider window is re-truncated on the way to the model, so it is clamped instead, with a RuntimeWarning when the caller is the one who asked for it. Options are what make the room small: on the English checkpoint a 2-option question leaves 483 tokens for the state and a 100-option one leaves 100. A smaller window isolates a localized signal better (a short deciding span is a larger fraction of its window, so that window classifies it clearly), at the cost of more windows; the large default favors context and throughput. noul is robust to this, choice/score benefit from a smaller window when the deciding span is a small part of a long, otherwise-neutral document.
stride (Optional[int], default: None ) –

token step between windows. Defaults to half the effective window (50% overlap), so a span near a boundary still lands whole inside some window. A stride past the effective window is refused rather than clamped: the tokens between each pair of windows would be read by no window at all, which is the failure this method exists to prevent.
aggregate (str, default: 'auto' ) –

"auto" (the per-type rules above) is the only mode for now.
batch_size (Optional[int], default: None ) –

cap on windows per forward pass, to bound memory on very long states.
lang (Optional[str], default: None ) –

per-language temperature selection, as in system_one.
hooks (HookArg, default: None ) –

Per-call hooks, appended after any installed on the Agent. See laya.hooks.
on_predict_start (PredictHookArg, default: None ) –

A per-call start hook, as in system_one.
on_predict_end (PredictHookArg, default: None ) –

A per-call end hook, as in system_one.
hooks_raise (Optional[bool], default: None ) –

Override the Agent's hooks_raise for this call.
hooks_timeout (Optional[float], default: None ) –

Override the Agent's hooks_timeout for this call.

Raises:

ValueError –

aggregate is anything but "auto"; the questions' options fill the whole sequence, leaving no room for the state; or stride steps past the effective window, so tokens between two windows would be read by nothing.

Returns a single result dict, the same shape as system_one, with usage["windows"] added. The key is always present and counts the windows the model scored to produce the answer: 1 for a state that fit one window, N for a document scanned in N overlapping windows (or the N a start hook rewrote them to), and 0 when a start hook answered the document, or left no states to score, before any window was read -- on either path, so a cached answer never reads as a window the model read.

Across several windows the truncation keys are combined like every other usage field: truncated, state_tokens and state_tokens_dropped are summed (so truncated is the number of windows that were cut, and the token counts include the overlap), and truncated_questions is the last window's list. The two can disagree: when only an earlier window was cut, truncated is above 0 and truncated_questions is empty. A window is cut when it is larger than the room a question's head leaves. Test usage["truncated"] > 0 here, not is True.

system_one ¶

system_one(
    state: Union[str, dict, list],
    questions: Dict[str, Dict[str, Any]],
    lang: Optional[str] = None,
    hooks=None,
    on_predict_start=None,
    on_predict_end=None,
    hooks_raise: Optional[bool] = None,
    hooks_timeout: Optional[float] = None,
    max_len: Optional[int] = None,
    head_max_len: Optional[int] = None,
    min_confidence: Optional[float] = None,
) -> Dict[str, Any]

Evaluate typed questions across state in a single, parallel forward pass.

Parameters:

state (Union[str, dict, list]) –

Text string, JSON dict, or conversation turn list.
questions (Dict[str, Dict[str, Any]]) –

Dictionary mapping question_id -> question definition. - choice: {"type": "choice", "instructions": "...", "criteria": {"optA": "...", ...}} - score: {"type": "score", "instructions": "...", "criteria": ["lvl0", "lvl1", ...]} - noul: {"type": "noul", "instructions": "...", "criteria": {"false": "...", "true": "..."}, "labels": {"false": "B", "true": "A"}}

Noul criteria and labels are optional. Labels only control the text shown to the model; their keys retain false/true semantics, and the returned noul value is always P(true). Labels default to false/true for compatibility.

Returns:

Dict[str, Any] –

Dictionary with answers, probabilities, calibrated confidence, and token usage. Empty questions return empty answers and zero token usage without tokenization or a model forward pass.

When the head budget leaves two options with the same token span, usage carries an options entry for each question it happened to -- total, distinct and tokens_per_option -- because an answer chosen among 42 distinguishable spans of 58 has a ceiling that is the budget's and not the model's. Questions whose options all survive are absent, so a request that collapses nothing is unchanged.

usage also reports whether the state fit: truncated, state_tokens, state_tokens_dropped, and truncated_questions (the questions whose head left too little room). A caller that cares whether the answer saw the whole state should read usage["truncated"] rather than estimate from the length of what it sent.

To score many states at once, see predict_batch, which shares forward passes across them.

decide ¶

decide(
    state: Union[str, dict, list],
    schema: Any = None,
    *,
    questions: Optional[Dict[str, Any]] = None,
    return_details: bool = False,
    min_confidence: Optional[float] = None,
    **predict_kwargs,
) -> Any

Answer state against a schema (JSON schema or pydantic model) and return typed values.

See laya.structured. Pass exactly one of schema or questions; extra keyword arguments are forwarded to predict / system_one.

decide_batch ¶

decide_batch(
    states: List[Union[str, dict, list]],
    schema: Any = None,
    *,
    questions: Optional[Dict[str, Any]] = None,
    return_details: bool = False,
    min_confidence: Optional[float] = None,
    **predict_kwargs,
) -> List[Any]

Answer many states against one schema (JSON schema or pydantic model) in one batched call.

The throughput form of :meth:decide: the schema is planned once and its questions run over every state through :meth:predict_batch (shared forward passes, results in input order), then each state's answers are projected as decide does. Extra keyword arguments (batch_size=, lang=, hooks=, ...) are forwarded to predict_batch. See laya.structured.

fit_temperatures ¶

fit_temperatures(
    records, compute_ece: bool = False, seed: int = 0
) -> Dict[str, Any]

Fit per-bucket temperatures from CPU records and store them on this agent.

records are (qtype, logits, target, k). Build them with laya.calibrate.records_from_labeled when you have labeled forwards; this method does not download weights or write model.safetensors. seed only affects the held-out ECE split when compute_ece is true. The checkpoint cfg is left as loaded.

fit_binning ¶

fit_binning(
    records, min_bucket_n: int = MIN_BINNING_BUCKET_N
) -> Dict[str, Any]

Fit a histogram-binning map on top of this agent's fitted temperatures and store it.

records are the same (qtype, logits, target[, k]) tuples as fit_temperatures consumed. The map is keyed exactly like temperature_by_options, composes on top of the current temperatures, and is written out by save_calibration as binning_map.

save_calibration ¶

save_calibration(path: str) -> None

Write temperature, temperature_by_options, binning_map when the agent has one, and the checkpoint they were fitted for. Does not write weights.

load_calibration ¶

load_calibration(path: str) -> None

Read a JSON map written by save_calibration onto this agent.

The file carries three fields and all three are installed: temperature, temperature_by_options and binning_map. A file with no binning_map key installs None, so loading one clears a map this agent's fit_binning fitted -- the file is the whole calibration state, not a patch onto the current one.

A file with no version is treated as version 1 and still loads. A newer file whose recorded checkpoint does not match this agent warns and still loads. Temperatures that are not numbers, or that sit outside [TEMP_MIN, TEMP_MAX], are clamped with clamp_temperature the same way checkpoint load is; binning values are not clamped, they are refused with a ValueError naming the field, because an out-of-range binning value would move a confidence with nothing to fall back to.

load ¶

load(
    model_id_or_path: str = "convaiinnovations/laya",
    device: Optional[str] = None,
    token: Optional[str] = None,
    subfolder: Optional[str] = None,
    fast: bool = False,
    compile: bool = False,
    revision: Optional[str] = None,
    expected_sha256: Optional[Dict[str, str]] = None,
    lang_temperatures: Optional[Dict[str, Dict[str, Any]]] = None,
    hooks=None,
    on_predict_start=None,
    on_predict_end=None,
    hooks_raise: bool = True,
    hooks_concurrent: bool = True,
    hooks_timeout: Optional[float] = None,
    calibration: Optional[str] = None,
    backend: Optional[str] = None,
    onnx_path: Optional[str] = None,
    compile_warmup: bool = True,
    compile_cache: bool = False,
    compile_mode: str = "default",
) -> Union[Agent, ONNXAgent]

Load a Laya agent.

subfolder picks one checkpoint out of a repo that bundles several:

laya.load("convaiinnovations/laya")                           # English (repo root)
laya.load("convaiinnovations/laya", subfolder="multilingual")
laya.load("convaiinnovations/laya", fast=True)                # TileLang GPU fast path
laya.load("convaiinnovations/laya", compile=True)              # torch.compile the model

model_id_or_path also accepts a checkpoint name or alias -- the same ones Router resolves, so both entry points read one table:

laya.load("typed-decisions")
laya.load("ml")                                               # multilingual

Anything else (a Hub repo id, a local directory) is passed to Agent unchanged.

backend selects "auto", "eager", "compile", "tilelang" or "onnx". ONNX returns the existing ONNXAgent, with onnx_path (default "laya.onnx"). Other backends use Agent; an explicit backend takes precedence over the legacy flags.

revision/expected_sha256 pin and verify the downloaded artifacts; see Agent. hooks / on_predict_start / on_predict_end observe or shape every prediction; see laya.hooks. calibration is the same optional JSON path accepted by Agent.

ONNXAgent ¶

ONNXAgent(
    model_id_or_path: str,
    onnx_path: str = "laya.onnx",
    token: Optional[str] = None,
    subfolder: Optional[str] = None,
    revision: Optional[str] = None,
    expected_sha256: Optional[Dict[str, str]] = None,
    hooks=None,
    on_predict_start=None,
    on_predict_end=None,
    hooks_raise: bool = True,
    hooks_concurrent: bool = True,
    hooks_timeout: Optional[float] = None,
    lang_temperatures: Optional[Dict[str, Dict[str, Any]]] = None,
    calibration: Optional[str] = None,
)

Bases: HookRegistry

System 1 decision model runtime via ONNX: fast CPU-optimized decisions.

Load a Laya agent backed by ONNX Runtime.

Parameters:

model_id_or_path (str) –

HuggingFace Hub ID or local path to the original PyTorch checkpoint (used to load the tokenizer and config).
onnx_path (str, default: 'laya.onnx' ) –

Path to the exported .onnx file.
token (Optional[str], default: None ) –

Optional HuggingFace token for a private or gated checkpoint; falls back to $HF_TOKEN, exactly as Agent does. Only the tokenizer and config are fetched -- the graph itself is the local onnx_path.
subfolder (Optional[str], default: None ) –

Optional subfolder if downloading from a repo bundle.
revision (Optional[str], default: None ) –

Optional Hub revision (commit SHA/branch/tag). When omitted, huggingface_hub's normal default and existing offline cache are used.
expected_sha256 (Optional[Dict[str, str]], default: None ) –

Optional {path relative to the checkpoint dir: hexdigest} verified before any checkpoint file is parsed; opt-in, and applies to local directories too. A missing artifact raises FileNotFoundError and a digest mismatch raises ValueError; either error refuses the load.
hooks (HookArg, default: None ) –

Opt-in prediction hooks; see laya.hooks.
on_predict_start (PredictHookArg, default: None ) –

An opt-in start hook, run before inference.
on_predict_end (PredictHookArg, default: None ) –

An opt-in end hook, run after inference.
hooks_raise (bool, default: True ) –

When False, a failing hook warns and inference continues.
hooks_concurrent (bool, default: True ) –

When False, hooks are serialised with a lock.
hooks_timeout (Optional[float], default: None ) –

Bounds each hook call in seconds; None means no limit.
lang_temperatures (Optional[Dict[str, Dict[str, Any]]], default: None ) –

Optional per-language temperature overrides, keyed by language code, each {"temperature": [3 floats], "temperature_by_options": {}}. Applied when a lang= is passed to system_one/predict, mirroring the PyTorch Agent; a cross-backend swap otherwise loses calibration.

load_calibration ¶

load_calibration(path: str) -> None

Read a JSON map written by save_calibration onto this agent.

system_one ¶

system_one(
    state: Union[str, dict, list],
    questions: Dict[str, Dict[str, Any]],
    lang: Optional[str] = None,
    hooks=None,
    on_predict_start=None,
    on_predict_end=None,
    hooks_raise: Optional[bool] = None,
    hooks_timeout: Optional[float] = None,
    max_len: Optional[int] = None,
    head_max_len: Optional[int] = None,
    min_confidence: Optional[float] = None,
) -> Dict[str, Any]

Evaluate typed questions across state in one ONNX Runtime session run.

lang selects a per-language temperature override (see lang_temperatures), matching the PyTorch Agent.system_one signature so either backend is a drop-in for the other.

Defined in terms of predict_batch, exactly as the PyTorch Agent.system_one is, so the single-state and batched paths cannot drift apart.

Parameters:

state (Union[str, dict, list]) –

Text string, JSON dict, or conversation turn list.
questions (Dict[str, Dict[str, Any]]) –

Question definitions, with the shapes Agent.system_one accepts.
lang (Optional[str], default: None ) –

Per-language temperature override (see lang_temperatures).
hooks (HookArg, default: None ) –

Per-call hooks, appended after any installed on the agent.
on_predict_start (PredictHookArg, default: None ) –

A per-call start hook. It may rewrite the state/questions or call ctx.skip(...) to short-circuit inference.
on_predict_end (PredictHookArg, default: None ) –

A per-call end hook. It may rewrite the results.
hooks_raise (Optional[bool], default: None ) –

Override the agent's hooks_raise for this call.
hooks_timeout (Optional[float], default: None ) –

Override the agent's hooks_timeout for this call.
max_len (Optional[int], default: None ) –

Override the config's max_len for this call.
head_max_len (Optional[int], default: None ) –

Override the config's head_max_len for this call.
min_confidence (Optional[float], default: None ) –

Opt-in abstention threshold on answer_confidence (#361); an answer below it is returned flagged with low_confidence: True.

Returns:

Dict[str, Any] –

Dictionary with answers, probabilities, calibrated confidence, and token usage.

To score many states at once, see predict_batch, which shares session runs across them.

predict_batch ¶

predict_batch(
    states: List[Union[str, dict, list]],
    questions: Dict[str, Dict[str, Any]],
    batch_size: Optional[int] = None,
    lang: Optional[str] = None,
    hooks=None,
    on_predict_start=None,
    on_predict_end=None,
    hooks_raise: Optional[bool] = None,
    hooks_timeout: Optional[float] = None,
    max_len: Optional[int] = None,
    head_max_len: Optional[int] = None,
    sort_by_length: bool = False,
    min_confidence: Optional[float] = None,
) -> List[Dict[str, Any]]

Evaluate the same questions over many states, sharing ONNX Runtime session runs.

The throughput path, mirroring laya.agent.Agent.predict_batch: system_one collates one state's question rows per session run, so N states cost N runs. predict_batch collates several states' rows into one run -- or ceil(len(states) / batch_size) of them -- which is where ONNX Runtime's own parallelism pays off on CPU.

Parameters:

states (List[Union[str, dict, list]]) –

A list of states (each a text string, JSON dict, or conversation turn list). The same questions are evaluated against every state.
questions (Dict[str, Dict[str, Any]]) –

Question definitions, exactly as accepted by system_one.
batch_size (Optional[int], default: None ) –

Optional cap on states per session run. None sends them all in one run; set it to bound peak memory when batching many or long states.
lang (Optional[str], default: None ) –

Per-language temperature override applied to every state; see lang_temperatures.
hooks (HookArg, default: None ) –

Per-call hooks, appended after any installed on the agent.
on_predict_start (PredictHookArg, default: None ) –

A per-call start hook. It may rewrite the states/questions or call ctx.skip(...) to short-circuit inference.
on_predict_end (PredictHookArg, default: None ) –

A per-call end hook. It may rewrite the results.
hooks_raise (Optional[bool], default: None ) –

Override the agent's hooks_raise for this call.
hooks_timeout (Optional[float], default: None ) –

Override the agent's hooks_timeout for this call.
max_len (Optional[int], default: None ) –

Override the config's max_len for this call.
head_max_len (Optional[int], default: None ) –

Override the config's head_max_len for this call.
sort_by_length (bool, default: False ) –

Group similarly sized encoded states within windows of eight batches to reduce padding, exactly as Agent.predict_batch does. Requires an explicit batch_size greater than one and smaller than the number of states; otherwise it has no effect. Results retain input order. Changing batch shapes can slightly change floating-point predictions near decision thresholds.
min_confidence (Optional[float], default: None ) –

Opt-in abstention threshold on answer_confidence (#361); answers below it are returned flagged with low_confidence: True.

Returns:

List[Dict[str, Any]] –

A list of per-state result dicts, each identical in shape to system_one's output and aligned with states by index.

predict_long ¶

predict_long(
    state: Union[str, dict, list],
    questions: Dict[str, Dict[str, Any]],
    window: Optional[int] = None,
    stride: Optional[int] = None,
    aggregate: str = "auto",
    batch_size: Optional[int] = None,
    lang: Optional[str] = None,
    hooks=None,
    on_predict_start=None,
    on_predict_end=None,
    hooks_raise: Optional[bool] = None,
    hooks_timeout: Optional[float] = None,
) -> Dict[str, Any]

Evaluate questions over a state longer than the context window, scanning it in overlapping windows and aggregating per question.

The ONNX port of laya.agent.Agent.predict_long, with the same aggregation rules: system_one truncates a state that exceeds max_len to a single window, silently dropping the rest. predict_long tokenizes the state once, splits it into overlapping token windows, scores every window through predict_batch -- so the windows share ONNX Runtime session runs rather than costing one each -- and combines the per-window answers:

noul -> P(true) is the max over windows (the statement holds if any window supports it)
choice-> the answer from the single most-confident window, so a localized signal isn't out-voted by the many neutral windows a long document is mostly made of
score -> the level from the most-confident window, likewise

The returned probability/confidence is the deciding window's, not a calibrated number for the whole document, for the same reasons the PyTorch docstring gives. Each answer carries answer["window"] -- the deciding window's index, token_start/token_end into the tokenized state, and the window count.

A state the questions leave room for is passed straight to system_one (identical output), since system_one reads it whole; with an explicit window, a state that fits that window.

Parameters:

state (Union[str, dict, list]) –

Text string, JSON dict, or conversation turn list.
questions (Dict[str, Dict[str, Any]]) –

Question definitions, exactly as accepted by system_one.
window (Optional[int], default: None ) –

State tokens per window. Defaults to the per-question state budget (max(64, max_len - head_max_len - 8); the 64 is a floor). Smaller windows isolate a localized signal better at the cost of more windows, as in Agent.predict_long.
stride (Optional[int], default: None ) –

Token step between windows. Defaults to half the effective window -- the window these questions leave room for, which is smaller than the window that was asked for once the room caps it (50% overlap either way).
aggregate (str, default: 'auto' ) –

"auto" (the per-type rules above) is the only mode for now.
batch_size (Optional[int], default: None ) –

Cap on windows per session run, to bound peak memory on very long states.
lang (Optional[str], default: None ) –

Per-language temperature selection, as in system_one.
hooks (HookArg, default: None ) –

Per-call hooks, appended after any installed on the agent. They follow Agent.predict_long's contract: they wrap the inference that answers the state, a start hook that answers with ctx.skip(...) gets usage["windows"] == 0 and no window attribution, and a rewritten scan is aggregated without answer["window"].
on_predict_start (PredictHookArg, default: None ) –

A per-call start hook, as in system_one.
on_predict_end (PredictHookArg, default: None ) –

A per-call end hook, as in system_one.
hooks_raise (Optional[bool], default: None ) –

Override the agent's hooks_raise for this call.
hooks_timeout (Optional[float], default: None ) –

Override the agent's hooks_timeout for this call.

Returns a single result dict, the same shape as system_one, with usage["windows"] added. Across several windows the truncation keys are summed or carried the same way as in Agent.predict_long: truncated is a window count and truncated_questions is the last window's list, so truncated can be above 0 while the list is empty.

decide ¶

decide(
    state: Union[str, dict, list],
    schema: Any = None,
    *,
    questions: Optional[Dict[str, Dict[str, Any]]] = None,
    return_details: bool = False,
    min_confidence: Optional[float] = None,
    **predict_kwargs,
) -> Any

Answer state against a schema (JSON schema or pydantic model) and return typed values.

See laya.structured. Pass exactly one of schema or questions; extra keyword arguments are forwarded to predict / system_one.

decide_batch ¶

decide_batch(
    states: List[Union[str, dict, list]],
    schema: Any = None,
    *,
    questions: Optional[Dict[str, Dict[str, Any]]] = None,
    return_details: bool = False,
    min_confidence: Optional[float] = None,
    **predict_kwargs,
) -> List[Any]

Answer many states against one schema through predict_batch; see laya.structured.

Quantized export¶

scripts/export_onnx.py --quantize writes an INT8 weight-only quantized copy beside the fp32 export (laya.onnx also produces laya.int8.onnx). Dynamic quantization converts the MatMul weights to int8 with the activation scale computed per input at run time, so no calibration dataset is needed, and ONNXAgent loads the result by pointing onnx_path at it. On CPU it is roughly 2x faster than the eager model and ~1.8x faster than the fp32 ONNX graph, and 1.4-2.8x smaller depending on the checkpoint.

INT8 trades real accuracy, so it is a size/latency option, not a free one — do not use it where the calibrated probability or confidence matters. Scales are per-tensor by default; --per-channel opts into per-channel weights but on the dynamic path that collapses the decision model (agreement with the eager model dropped to ~32% on the English checkpoint and ~40% on the multilingual one, vs ~67% / ~83% per-tensor; see issue #790). Even per-tensor drifts noticeably on the larger checkpoint; accuracy-safe int8 would need QAT or SmoothQuant-style outlier handling. The int8 graph is CPU-only: ONNX Runtime has no INT8 MatMul kernel on the CUDAExecutionProvider, and a GPU provider silently falls back per node.

python scripts/export_onnx.py --model convaiinnovations/laya --output laya.onnx --quantize

In-process dynamic quantization¶

The quantized export above is the int8 path ONNXAgent loads. torch.ao.quantization.quantize_dynamic on a live Agent is a different operation.

Issue #1065 reports that the default fbgemm engine crashes with SIGILL (illegal instruction) on a CPU with SSE4.2 and no AVX2 (laya 0.3.22, torch 2.14). Check the CPU first. torch.backends.cpu.get_cpu_capability() returns "NO AVX" when AVX2 is unavailable, and "AVX2" or "AVX512" when it is. On Linux, grep -m1 flags /proc/cpuinfo prints the raw flags (avx2 is absent in the report). The setting that ran there is qnnpack, applied only to the encoder:

import torch
import torch.nn as nn

torch.backends.quantized.engine = "qnnpack"
torch.ao.quantization.quantize_dynamic(
    agent.model.encoder, {nn.Linear}, dtype=torch.qint8, inplace=True)

Encoder-only quantization does run, with either engine the CPU can execute. It is not a speedup you can assume, and it changes scores. On the machine in #1065, qnnpack took 3.7 s against 3.0 s in fp32 for 20 states, and one candidate's noul went from 0.15 to 0.04. A spot check on the multilingual checkpoint with torch 2.14.1 (capability AVX512, four threads, short states) saw the same kind of move: qnnpack took 2.1 s against 0.57 s, and one noul went from 0.90 to 0.13. fbgemm was faster on that CPU (0.39 s) and still moved the score, to 0.15. Validate latency and probabilities on your own data.

Do not quantize the whole agent.model. That replaces the decision head's nn.Linear layers with dynamically quantized linears, whose weight is a method rather than a tensor, and two different reads of that method can be the one that raises.

On every published checkpoint it is the encoder stack. agent.model.head is an nn.TransformerEncoder of head_layers nn.TransformerEncoderLayers (laya/common.py::DecisionModel.__init__, and None when head_layers is 0), and each layer's sparsity fast path reads .device on the quantized weights, so predict raises AttributeError: 'function' object has no attribute 'device'. Reproduced on laya, laya-multilingual and laya-typed-decisions, qnnpack engine, from torch/nn/modules/transformer.py's (x.device.type in _supported_device_type) line.

Whether that fast path is taken is what decides which message you get, so a reader who sees the other one has not found a different bug. Torch skips the fast path when self_attn.num_heads is odd, and the head takes nhead = max(1, d // 64), so a model built from a config whose hidden size is under 128 gets nhead = 1, survives the stack, and fails one step later in DecisionModel.forward, which reads .dtype on act_head[0].weight whenever autocast is off: AttributeError: 'function' object has no attribute 'dtype'. The published checkpoints are d=768 and d=1024, so nhead is 12 or 16 and never that case.

Either way, leave the head in fp32.

Agent¶

Agent ¶

backend property ¶

backend_object property ¶

set_backend ¶

accelerate ¶

warmup ¶

deaccelerate ¶

dtype_for ¶

predict_batch ¶

predict_long ¶

system_one ¶

decide ¶

decide_batch ¶

fit_temperatures ¶

fit_binning ¶

save_calibration ¶

load_calibration ¶

load ¶

ONNXAgent ¶

load_calibration ¶

system_one ¶

predict_batch ¶

predict_long ¶

decide ¶

decide_batch ¶

Quantized export¶

In-process dynamic quantization¶

backend `property` ¶

backend_object `property` ¶