tjameswilliams/laya-r-mlx

Rust inference engine for Laya decision models on Apple silicon (MLX, candle/Metal, C ABI, Swift, runs on iPhone)

Rust

0

1 commits

updated Sep 23, 2026

See the code

README

laya-r-mlx

License: Apache 2.0 Platform: macOS · iOS (Apple silicon) Rust 1.85+ Model: convaiinnovations/laya

laya-r-mlx is a native, Python-free inference engine for the Laya System 1 decision models. Laya answers typed questions (choice, score, noul) about any state (text, e-mail, ticket, JSON, conversation) in a single encoder forward pass, with calibrated probabilities and no text generation. This engine gives the same answers as the Python package, faster, from one static library you can drop into a macOS app, an iPhone app, a CLI, or anything with a C FFI.

  • Same answers. Tokenization, prompt layout, calibration and the result JSON are verified byte-identical to laya 0.3.10 against recorded fixtures. In f16 every choice matches and probabilities agree to 0.002; in f32 the result JSON is byte-identical.
  • Two backends, one trait. MLX (via mlx-rs) for the fastest Mac path, and candle (Metal or CPU, pure Rust) which also builds for iOS.
  • Embeddable. A C ABI (laya.h), an XCFramework, a Swift package, and a SwiftUI demo app that runs on an iPhone.

Inference numbers

One predict call end to end (tokenize, forward, decode), p50 of 30 calls after warm-up. "E-mail, 4 questions" is the example below (348 tokens across the four question rows); "1 question" is a single noul question on the same e-mail. Full tables, methodology and caveats in BENCHMARKS.md.

MacBook Pro, M4 Max (Python: laya 0.3.10 on torch 2.14, MPS, f32, same machine and inputs)

checkpointinputPython (torch MPS)Rust MLXRust candle (Metal)
laya (ModernBERT-large, 421M)e-mail, 4 questions37.2 ms32.2 ms32.4 ms
1 question18.7 ms15.4 ms14.8 ms
512-token state, 4 questions170 ms173 ms162 ms
laya-multilingual (mmBERT-base, 322M)e-mail, 4 questions18.0 ms13.0 ms16–21 ms
1 question11.1 ms7.3 ms12.3 ms
model load~2 s + torch import0.2–0.4 s0.2–0.5 s

iPhone 17 Pro (candle backend, Metal f16, laya-multilingual bundled in the demo app)

iPhone 17 Pro
e-mail, 4 questions66 ms p50 (64 ms min)
1 question20 ms p50
load 644 MB checkpoint0.7 s (2.4 s on first launch)
warm-up (Metal pipelines)25 ms warm cache, ~1.3 s cold

The forward pass is GPU-compute-bound at these sizes (about 22 ms of the 32 ms is f16 matrix multiplication at the M4 Max's peak), so the gain over torch on the Mac is framework overhead and load time; on the phone the point is that it runs at all, at interactive latency.

Quick start

Checkpoints are read from a directory or from your local Hugging Face cache (the engine does not download). Get them once with the hf CLI (or by running the Python package):

hf download convaiinnovations/laya --include "rl_agent_config.json" "model.safetensors" "tokenizer/*" "encoder/*"
hf download convaiinnovations/laya --include "multilingual/*"      # optional, 100+ languages

Build and ask a question:

cargo build --release -p laya-cli      # first build compiles MLX from source (~5 min)
target/release/laya predict --pretty \
  --state '{"from": "user@acme.com", "subject": "Duplicate charge on invoice #4411",
            "body": "Hi, we were billed twice for March. Please refund the duplicate today or we will cancel our plan."}' \
  --questions '{
    "department": {"type": "choice", "instructions": "Which department should handle this request?",
                   "criteria": {"billing": "invoices, payments, refunds", "technical": "bugs, outages, system errors",
                                "sales": "pricing, new contracts", "other": "everything else"}},
    "urgency":    {"type": "score", "instructions": "How urgent is this request?",
                   "criteria": ["not urgent", "soon", "critical deadline or blocking issue"]},
    "churn_risk": {"type": "noul", "instructions": "Does the user threaten to cancel or leave?"}}'
{
  "model": "laya-rl-agent",
  "answers": {
    "department": {"type": "choice", "choice": "billing",
                   "probabilities": {"billing": 0.9653, "technical": 0.014, "sales": 0.0101, "other": 0.0107},
                   "confidence": 0.864, "action": {"act_probability": 1.0}},
    "urgency":    {"type": "score", "score": 1.44, "legend": {"0": "not urgent", "1": "soon", "2": "critical deadline or blocking issue"},
                   "probabilities": {"0": 0.1164, "1": 0.3271, "2": 0.5565}, "confidence": 0.1425, "action": {"act_probability": 1.0}},
    "churn_risk": {"type": "noul", "noul": 0.8248, "confidence": 0.8248, "action": {"act_probability": 1.0}}
  },
  "usage": {"input_tokens": 264, "output_tokens": 0}
}

The result has exactly the Python package's shape, so anything written against laya's predict() output works unchanged.

How to use

The three question types

typecriteriaanswer
choice{label: description} (or a list of labels)choice (the label) + probabilities per label
scorelist of level descriptions, index 0 firstscore (expected level, float) + probabilities per level
nouloptional {"true": ..., "false": ...}; optional labelsnoul = P(true)

Every answer also carries confidence (normalized-entropy, temperature-calibrated by the checkpoint) and action.act_probability. A state can be a plain string, a JSON object, or a JSON array of conversation turns (arrays are truncated from the left so the newest turn survives). Key order matters: objects are serialized in key order before tokenization and a choice question's criteria order is the option order, so keep both stable if you want the same answer on every platform.

Command line

laya predict --state <text | JSON | @file> --questions <JSON | @file> [--pretty]
laya bench   [--iters 30] [--case email_main] [--batch 8]     # p50/mean/min + encode/forward/decode split
laya parity                                                  # compare with the Python fixtures
laya info                                                    # resolved checkpoint and its config

Common options on every subcommand:

flagmeaning
--model <hub id or dir>default convaiinnovations/laya
--subfolder multilingual | typed-decisionspick a bundled checkpoint
--backend mlx | candle | autoauto prefers MLX when compiled in
--device auto | gpu | cpu
--f32compute in f32: byte-identical to Python, same speed on M-series before M5, 2× memory

Build without a backend you don't want: cargo build --release -p laya-cli --no-default-features --features candle.

Rust

[dependencies]
laya-core = { git = "https://github.com/TJamesWilliams/laya-r-mlx" }
laya-mlx  = { git = "https://github.com/TJamesWilliams/laya-r-mlx" }   # or laya-candle
use laya_core::{Agent, BackendOptions, resolve::resolve_model_dir};
use serde_json::json;

let dir = resolve_model_dir("convaiinnovations/laya", None)?;          // or a local path
let agent = Agent::load(&dir, &BackendOptions::default(), Box::new(laya_mlx::make_backend))?;

let result = agent.predict(
    &json!({"body": "We were billed twice, refund us today or we cancel."}),
    &json!({"refund": {"type": "noul", "instructions": "Does the user request a refund?"}}),
)?;
println!("{}", result["answers"]["refund"]["noul"]);

// Many states, same questions, shared forward passes:
let results = agent.predict_batch(&states, &questions, Some(16))?;

Agent is Send + Sync; share one per checkpoint across threads. laya-core enables serde_json's preserve_order feature, so json! objects keep their key order. predict_batch_timed returns the encode/forward/decode split for profiling.

C and other languages

cargo build --release -p laya-ffi produces target/release/liblaya.a and liblaya.dylib; the header is crates/laya-ffi/include/laya.h.

char *err = NULL;
laya_agent *a = laya_load("/path/to/checkpoint", NULL, "auto", LAYA_DEVICE_AUTO, 0, &err);
laya_warmup(a, &err);
char *json = laya_predict(a, state_json, questions_json, &err);   /* same JSON as the CLI */
laya_string_free(json);
laya_agent_free(a);

A complete example with error handling and timing is in crates/laya-ffi/examples/smoke.c; link with -framework Metal -framework Foundation -framework Accelerate -lc++.

Swift, macOS and iOS

scripts/build-xcframework.sh          # dist/Laya.xcframework: macOS arm64, iOS arm64, iOS simulator

Add swift/LayaKit as a local package (it references dist/Laya.xcframework), then:

import LayaKit

let agent = try LayaAgent(modelPath: modelDir.path)    // folder with rl_agent_config.json, model.safetensors, tokenizer/, encoder/
try agent.warmUp()                                      // once, off the main thread: compiles Metal pipelines
let result = try agent.predict(
    stateJSON: #"{"body": "Refund me now or I cancel"}"#,
    questionsJSON: #"{"refund": {"type": "noul", "instructions": "Refund requested?"}}"#)
let p = (result["answers"] as? [String: Any])?["refund"]   // {"noul": 0.97, ...}

Prefer the JSON-text entry points over Swift dictionaries: [String: Any] does not preserve key order, which changes the input the model sees. The XCFramework contains the candle backend (MLX's C++ objects cannot be packaged in an XCFramework; macOS apps that want MLX link liblaya.a directly as in the C example).

Demo app. swift/LayaDemo is a SwiftUI app for iOS and macOS that loads a checkpoint from its bundle or from Documents/Model, answers the example questions and times 20 calls:

scripts/prepare-demo-model.sh multilingual   # copies the checkpoint into swift/LayaDemo/Model (647 MB)
cd swift/LayaDemo && xcodegen generate && open LayaDemo.xcodeproj   # pick your iPhone, run

For iOS builds use the rustup toolchain (rustup target add aarch64-apple-ios aarch64-apple-ios-sim); Homebrew's rustc has no iOS targets. The CLI itself also runs inside the iOS simulator:

~/.cargo/bin/cargo build --release -p laya-cli --no-default-features --features candle --target aarch64-apple-ios-sim
xcrun simctl spawn booted target/aarch64-apple-ios-sim/release/laya bench --model <checkpoint dir>

Shipping tip: don't bundle the 644–842 MB checkpoint in the app binary; download it into Application Support on first run and pass that folder to LayaAgent.

Verifying against the Python package

uv venv && uv pip install laya                                   # any Python 3.10+
python scripts/ref_dump.py tests/fixtures/ref                     # answers + timings for both checkpoints
python scripts/ref_dump_head.py tests/fixtures/ref                # head-output vectors
cp tests/fixtures/ref/laya/fixtures.json tests/fixtures/laya.json
cp tests/fixtures/ref/multilingual/fixtures.json tests/fixtures/multilingual.json
cargo test --release --workspace                                  # core parity + per-backend parity
target/release/laya parity --backend candle --subfolder multilingual

The core tests assert byte-identical token ids, marker positions and result JSON; the backend tests assert every choice equal, logits within 0.15 and probabilities within 0.02 for f16.

Layout

crates/laya-core    tokenizer, sequence builder, decoding, weights, Backend trait (no ML framework)
crates/laya-mlx     MLX backend (macOS)
crates/laya-candle  candle backend (macOS Metal/CPU, iOS)
crates/laya-cli     `laya predict | bench | parity | info`
crates/laya-ffi     C ABI  ->  liblaya.a / liblaya.dylib + include/laya.h
swift/LayaKit       Swift package over dist/Laya.xcframework
swift/LayaDemo      SwiftUI demo app (xcodegen spec)
docs/MODEL.md       exact forward-pass spec the backends implement
scripts/            Python reference dump, XCFramework build, demo model prep
tests/fixtures      Python reference outputs (git-ignored; regenerate as above)

Not ported (yet)

  • The Python Router (script/language detection choosing a checkpoint per request), the e-mail cleaning presets, and the LangChain/MCP/HTTP integrations. Load the checkpoint you want.
  • Checkpoint download. Use hf download or the Python package once.
  • Linux/CUDA. laya-core is portable, but both backends target Apple GPUs; a CPU-only build of the candle backend works anywhere candle does.

License

Apache License 2.0, see LICENSE and NOTICE.

This engine reimplements the inference pipeline of Laya by Convai Innovations (Apache 2.0). The checkpoints it runs are published by Convai Innovations on Hugging Face under Apache 2.0 (their base encoders: ModernBERT, Apache 2.0; mmBERT, MIT) and are not part of this repository. All Rust dependencies are MIT and/or Apache 2.0 licensed.

apple-silicon
candle
inference
ios
laya
metal
mlx
rust

Contributors

tjameswilliams/laya-r-mlx

Rust inference engine for Laya decision models on Apple silicon (MLX, candle/Metal, C ABI, Swift, runs on iPhone)

Rust

0

1 commits

updated Sep 23, 2026

See the code

README

laya-r-mlx

License: Apache 2.0 Platform: macOS · iOS (Apple silicon) Rust 1.85+ Model: convaiinnovations/laya

laya-r-mlx is a native, Python-free inference engine for the Laya System 1 decision models. Laya answers typed questions (choice, score, noul) about any state (text, e-mail, ticket, JSON, conversation) in a single encoder forward pass, with calibrated probabilities and no text generation. This engine gives the same answers as the Python package, faster, from one static library you can drop into a macOS app, an iPhone app, a CLI, or anything with a C FFI.

  • Same answers. Tokenization, prompt layout, calibration and the result JSON are verified byte-identical to laya 0.3.10 against recorded fixtures. In f16 every choice matches and probabilities agree to 0.002; in f32 the result JSON is byte-identical.
  • Two backends, one trait. MLX (via mlx-rs) for the fastest Mac path, and candle (Metal or CPU, pure Rust) which also builds for iOS.
  • Embeddable. A C ABI (laya.h), an XCFramework, a Swift package, and a SwiftUI demo app that runs on an iPhone.

Inference numbers

One predict call end to end (tokenize, forward, decode), p50 of 30 calls after warm-up. "E-mail, 4 questions" is the example below (348 tokens across the four question rows); "1 question" is a single noul question on the same e-mail. Full tables, methodology and caveats in BENCHMARKS.md.

MacBook Pro, M4 Max (Python: laya 0.3.10 on torch 2.14, MPS, f32, same machine and inputs)

checkpointinputPython (torch MPS)Rust MLXRust candle (Metal)
laya (ModernBERT-large, 421M)e-mail, 4 questions37.2 ms32.2 ms32.4 ms
1 question18.7 ms15.4 ms14.8 ms
512-token state, 4 questions170 ms173 ms162 ms
laya-multilingual (mmBERT-base, 322M)e-mail, 4 questions18.0 ms13.0 ms16–21 ms
1 question11.1 ms7.3 ms12.3 ms
model load~2 s + torch import0.2–0.4 s0.2–0.5 s

iPhone 17 Pro (candle backend, Metal f16, laya-multilingual bundled in the demo app)

iPhone 17 Pro
e-mail, 4 questions66 ms p50 (64 ms min)
1 question20 ms p50
load 644 MB checkpoint0.7 s (2.4 s on first launch)
warm-up (Metal pipelines)25 ms warm cache, ~1.3 s cold

The forward pass is GPU-compute-bound at these sizes (about 22 ms of the 32 ms is f16 matrix multiplication at the M4 Max's peak), so the gain over torch on the Mac is framework overhead and load time; on the phone the point is that it runs at all, at interactive latency.

Quick start

Checkpoints are read from a directory or from your local Hugging Face cache (the engine does not download). Get them once with the hf CLI (or by running the Python package):

hf download convaiinnovations/laya --include "rl_agent_config.json" "model.safetensors" "tokenizer/*" "encoder/*"
hf download convaiinnovations/laya --include "multilingual/*"      # optional, 100+ languages

Build and ask a question:

cargo build --release -p laya-cli      # first build compiles MLX from source (~5 min)
target/release/laya predict --pretty \
  --state '{"from": "user@acme.com", "subject": "Duplicate charge on invoice #4411",
            "body": "Hi, we were billed twice for March. Please refund the duplicate today or we will cancel our plan."}' \
  --questions '{
    "department": {"type": "choice", "instructions": "Which department should handle this request?",
                   "criteria": {"billing": "invoices, payments, refunds", "technical": "bugs, outages, system errors",
                                "sales": "pricing, new contracts", "other": "everything else"}},
    "urgency":    {"type": "score", "instructions": "How urgent is this request?",
                   "criteria": ["not urgent", "soon", "critical deadline or blocking issue"]},
    "churn_risk": {"type": "noul", "instructions": "Does the user threaten to cancel or leave?"}}'
{
  "model": "laya-rl-agent",
  "answers": {
    "department": {"type": "choice", "choice": "billing",
                   "probabilities": {"billing": 0.9653, "technical": 0.014, "sales": 0.0101, "other": 0.0107},
                   "confidence": 0.864, "action": {"act_probability": 1.0}},
    "urgency":    {"type": "score", "score": 1.44, "legend": {"0": "not urgent", "1": "soon", "2": "critical deadline or blocking issue"},
                   "probabilities": {"0": 0.1164, "1": 0.3271, "2": 0.5565}, "confidence": 0.1425, "action": {"act_probability": 1.0}},
    "churn_risk": {"type": "noul", "noul": 0.8248, "confidence": 0.8248, "action": {"act_probability": 1.0}}
  },
  "usage": {"input_tokens": 264, "output_tokens": 0}
}

The result has exactly the Python package's shape, so anything written against laya's predict() output works unchanged.

How to use

The three question types

typecriteriaanswer
choice{label: description} (or a list of labels)choice (the label) + probabilities per label
scorelist of level descriptions, index 0 firstscore (expected level, float) + probabilities per level
nouloptional {"true": ..., "false": ...}; optional labelsnoul = P(true)

Every answer also carries confidence (normalized-entropy, temperature-calibrated by the checkpoint) and action.act_probability. A state can be a plain string, a JSON object, or a JSON array of conversation turns (arrays are truncated from the left so the newest turn survives). Key order matters: objects are serialized in key order before tokenization and a choice question's criteria order is the option order, so keep both stable if you want the same answer on every platform.

Command line

laya predict --state <text | JSON | @file> --questions <JSON | @file> [--pretty]
laya bench   [--iters 30] [--case email_main] [--batch 8]     # p50/mean/min + encode/forward/decode split
laya parity                                                  # compare with the Python fixtures
laya info                                                    # resolved checkpoint and its config

Common options on every subcommand:

flagmeaning
--model <hub id or dir>default convaiinnovations/laya
--subfolder multilingual | typed-decisionspick a bundled checkpoint
--backend mlx | candle | autoauto prefers MLX when compiled in
--device auto | gpu | cpu
--f32compute in f32: byte-identical to Python, same speed on M-series before M5, 2× memory

Build without a backend you don't want: cargo build --release -p laya-cli --no-default-features --features candle.

Rust

[dependencies]
laya-core = { git = "https://github.com/TJamesWilliams/laya-r-mlx" }
laya-mlx  = { git = "https://github.com/TJamesWilliams/laya-r-mlx" }   # or laya-candle
use laya_core::{Agent, BackendOptions, resolve::resolve_model_dir};
use serde_json::json;

let dir = resolve_model_dir("convaiinnovations/laya", None)?;          // or a local path
let agent = Agent::load(&dir, &BackendOptions::default(), Box::new(laya_mlx::make_backend))?;

let result = agent.predict(
    &json!({"body": "We were billed twice, refund us today or we cancel."}),
    &json!({"refund": {"type": "noul", "instructions": "Does the user request a refund?"}}),
)?;
println!("{}", result["answers"]["refund"]["noul"]);

// Many states, same questions, shared forward passes:
let results = agent.predict_batch(&states, &questions, Some(16))?;

Agent is Send + Sync; share one per checkpoint across threads. laya-core enables serde_json's preserve_order feature, so json! objects keep their key order. predict_batch_timed returns the encode/forward/decode split for profiling.

C and other languages

cargo build --release -p laya-ffi produces target/release/liblaya.a and liblaya.dylib; the header is crates/laya-ffi/include/laya.h.

char *err = NULL;
laya_agent *a = laya_load("/path/to/checkpoint", NULL, "auto", LAYA_DEVICE_AUTO, 0, &err);
laya_warmup(a, &err);
char *json = laya_predict(a, state_json, questions_json, &err);   /* same JSON as the CLI */
laya_string_free(json);
laya_agent_free(a);

A complete example with error handling and timing is in crates/laya-ffi/examples/smoke.c; link with -framework Metal -framework Foundation -framework Accelerate -lc++.

Swift, macOS and iOS

scripts/build-xcframework.sh          # dist/Laya.xcframework: macOS arm64, iOS arm64, iOS simulator

Add swift/LayaKit as a local package (it references dist/Laya.xcframework), then:

import LayaKit

let agent = try LayaAgent(modelPath: modelDir.path)    // folder with rl_agent_config.json, model.safetensors, tokenizer/, encoder/
try agent.warmUp()                                      // once, off the main thread: compiles Metal pipelines
let result = try agent.predict(
    stateJSON: #"{"body": "Refund me now or I cancel"}"#,
    questionsJSON: #"{"refund": {"type": "noul", "instructions": "Refund requested?"}}"#)
let p = (result["answers"] as? [String: Any])?["refund"]   // {"noul": 0.97, ...}

Prefer the JSON-text entry points over Swift dictionaries: [String: Any] does not preserve key order, which changes the input the model sees. The XCFramework contains the candle backend (MLX's C++ objects cannot be packaged in an XCFramework; macOS apps that want MLX link liblaya.a directly as in the C example).

Demo app. swift/LayaDemo is a SwiftUI app for iOS and macOS that loads a checkpoint from its bundle or from Documents/Model, answers the example questions and times 20 calls:

scripts/prepare-demo-model.sh multilingual   # copies the checkpoint into swift/LayaDemo/Model (647 MB)
cd swift/LayaDemo && xcodegen generate && open LayaDemo.xcodeproj   # pick your iPhone, run

For iOS builds use the rustup toolchain (rustup target add aarch64-apple-ios aarch64-apple-ios-sim); Homebrew's rustc has no iOS targets. The CLI itself also runs inside the iOS simulator:

~/.cargo/bin/cargo build --release -p laya-cli --no-default-features --features candle --target aarch64-apple-ios-sim
xcrun simctl spawn booted target/aarch64-apple-ios-sim/release/laya bench --model <checkpoint dir>

Shipping tip: don't bundle the 644–842 MB checkpoint in the app binary; download it into Application Support on first run and pass that folder to LayaAgent.

Verifying against the Python package

uv venv && uv pip install laya                                   # any Python 3.10+
python scripts/ref_dump.py tests/fixtures/ref                     # answers + timings for both checkpoints
python scripts/ref_dump_head.py tests/fixtures/ref                # head-output vectors
cp tests/fixtures/ref/laya/fixtures.json tests/fixtures/laya.json
cp tests/fixtures/ref/multilingual/fixtures.json tests/fixtures/multilingual.json
cargo test --release --workspace                                  # core parity + per-backend parity
target/release/laya parity --backend candle --subfolder multilingual

The core tests assert byte-identical token ids, marker positions and result JSON; the backend tests assert every choice equal, logits within 0.15 and probabilities within 0.02 for f16.

Layout

crates/laya-core    tokenizer, sequence builder, decoding, weights, Backend trait (no ML framework)
crates/laya-mlx     MLX backend (macOS)
crates/laya-candle  candle backend (macOS Metal/CPU, iOS)
crates/laya-cli     `laya predict | bench | parity | info`
crates/laya-ffi     C ABI  ->  liblaya.a / liblaya.dylib + include/laya.h
swift/LayaKit       Swift package over dist/Laya.xcframework
swift/LayaDemo      SwiftUI demo app (xcodegen spec)
docs/MODEL.md       exact forward-pass spec the backends implement
scripts/            Python reference dump, XCFramework build, demo model prep
tests/fixtures      Python reference outputs (git-ignored; regenerate as above)

Not ported (yet)

  • The Python Router (script/language detection choosing a checkpoint per request), the e-mail cleaning presets, and the LangChain/MCP/HTTP integrations. Load the checkpoint you want.
  • Checkpoint download. Use hf download or the Python package once.
  • Linux/CUDA. laya-core is portable, but both backends target Apple GPUs; a CPU-only build of the candle backend works anywhere candle does.

License

Apache License 2.0, see LICENSE and NOTICE.

This engine reimplements the inference pipeline of Laya by Convai Innovations (Apache 2.0). The checkpoints it runs are published by Convai Innovations on Hugging Face under Apache 2.0 (their base encoders: ModernBERT, Apache 2.0; mmBERT, MIT) and are not part of this repository. All Rust dependencies are MIT and/or Apache 2.0 licensed.

apple-silicon
candle
inference
ios
laya
metal
mlx
rust

Contributors

Languages

Rust

83.8%

Swift

7.7%

Python

4.5%

C

2.6%

Shell

1.4%