diff --git a/.agents/skills/README.md b/.agents/skills/README.md index fdf0bc9f94..4f9d99a0de 100644 --- a/.agents/skills/README.md +++ b/.agents/skills/README.md @@ -7,6 +7,6 @@ This directory contains specialized skills (recipes) to guide contributors and A - [Core Guidelines](./core-guidelines/SKILL.md) — Symmetrical architecture, conventions, and model rules. - [Add Native Extension](./add-native-extension/SKILL.md) — C++ operations and JSI bindings. - [Add Task Pipeline](./add-task-pipeline/SKILL.md) — TypeScript task pipelines and React hooks. -- [Model Schema Validation](./model-schema-validation/SKILL.md) — SymbolicTensor schemas and validation. +- [Model Schema Validation](./model-schema-validation/SKILL.md) — Model specs, dynamic shapes, and schema validation. - [Verify and Build](./verify-and-build/SKILL.md) — TypeScript typechecking, native rebuilding, and troubleshooting. - [Skills Maintenance](./skills-maintenance/SKILL.md) — Keeping skills synchronized with core primitives. diff --git a/.agents/skills/add-task-pipeline/SKILL.md b/.agents/skills/add-task-pipeline/SKILL.md index 45ec78d485..f8116dfa40 100644 --- a/.agents/skills/add-task-pipeline/SKILL.md +++ b/.agents/skills/add-task-pipeline/SKILL.md @@ -59,17 +59,17 @@ When implementing task constructors like `create` (e.g. `createClassifier` - Handle model-specific configuration parameters (such as unique normalization factors, thresholds, or label arrays) dynamically through the TypeScript task options argument rather than baking them rigidly into JSI C++ code or the model structure. This rule contrasts TypeScript options against values baked into C++ or the model; to choose between a TypeScript **option** and a TypeScript **constant**, see Principle 6. 6. **Options vs. Constants (bucket by who varies the value)**: - - Every parameter belongs in exactly one of three places. Decide by asking *who varies this, and when*: + - Every parameter belongs in exactly one of three places. Decide by asking _who varies this, and when_: - | The value... | Lives as | Example | - | ------------------------------------------------------------------------------------ | -------------------------------------------------------------- | -------------------------------------------------------------- | - | Varies across the shipped models/variants of the pipeline | A task **option**, set per-model in `models.ts` | Whisper `tiny`/`base`/`small` sizes; normalization factors; labels | - | Is fixed by the model architecture or the `.pte` export contract, identical for every shipped variant | A **`const`** in the task file, beside the code that reads it | Static tensor shapes; export-pinned scheduler/decoder scalars | - | Is a per-call choice made by the app developer | An **argument** to the worklet executor | `threshold`, `seed`, `prompt` | + | The value... | Lives as | Example | + | ----------------------------------------------------------------------------------------------------- | ------------------------------------------------------------- | ------------------------------------------------------------------ | + | Varies across the shipped models/variants of the pipeline | A task **option**, set per-model in `models.ts` | Whisper `tiny`/`base`/`small` sizes; normalization factors; labels | + | Is fixed by the model architecture or the `.pte` export contract, identical for every shipped variant | A **`const`** in the task file, beside the code that reads it | Static tensor shapes; export-pinned scheduler/decoder scalars | + | Is a per-call choice made by the app developer | An **argument** to the worklet executor | `threshold`, `seed`, `prompt` | - - **A parameter with exactly one valid value is not an option — it is a constant.** For single-model pipelines (e.g. `sdxsTextToImage`), exposing export-pinned scalars or static shapes as options advertises knobs that either fail schema validation or silently corrupt output when touched. Prefer a `const` with a comment naming *why* the value is fixed. - - **Corollary (a quick smell test):** if every variant in `models.ts` passes an *identical* options object, those fields are not configuration — move them into the task file as constants and shrink the model type to the paths that actually differ. - - Do not keep a loop, parameter, or code path solely because it *looks* more general. If the surrounding math is only valid for one value (e.g. coefficients pinned to a single distilled timestep), the generality is fake and the parameter is a correctness trap. + - **A parameter with exactly one valid value is not an option — it is a constant.** For single-model pipelines (e.g. `sdxsTextToImage`), exposing export-pinned scalars or static shapes as options advertises knobs that either fail schema validation or silently corrupt output when touched. Prefer a `const` with a comment naming _why_ the value is fixed. + - **Corollary (a quick smell test):** if every variant in `models.ts` passes an _identical_ options object, those fields are not configuration — move them into the task file as constants and shrink the model type to the paths that actually differ. + - Do not keep a loop, parameter, or code path solely because it _looks_ more general. If the surrounding math is only valid for one value (e.g. coefficients pinned to a single distilled timestep), the generality is fake and the parameter is a correctness trap. ## 🚫 Avoid / Anti-Patterns @@ -95,7 +95,7 @@ import type { WorkletRuntime } from 'react-native-worklets'; import { tensor } from '../../../core/tensor'; import { loadModel } from '../../../core/model'; -import { validateModelSchema, SymbolicTensor } from '../../../core/modelSchema'; +import { validateSpec, method, f32 } from '../../../core/schema'; import { wrapAsync } from '../../../core/runtime'; import { type ImageBuffer } from '../image'; import { createImagePreprocessor, type ImagePreprocessorOptions } from './preprocessing'; @@ -137,15 +137,15 @@ export async function createMyTask( const { modelPath, taskOpts } = config; const model = await wrapAsync(loadModel, runtime)(modelPath); - // Validate model schema - const meta = validateModelSchema( - model, - 'forward', - [SymbolicTensor('float32', [1, 3, 'H', 'W'], [3, 'H', 'W'])], - [SymbolicTensor('float32', [1, 10], [10])] - ); - const inpShape = meta.inputTensorMeta[0]!.shape; - const outShape = meta.outputTensorMeta[0]!.shape; + // Validate model spec + const { variant, dims } = validateSpec(model.schema, { + batched: method('forward', [f32(1, 3, 'H', 'W')], [f32(1, 10)]), + unbatched: method('forward', [f32(3, 'H', 'W')], [f32(10)]), + }); + + const [H, W] = dims.constant('H', 'W'); + const inpShape = { batched: [1, 3, H, W], unbatched: [3, H, W] }[variant]; + const outShape = { batched: [1, 10], unbatched: [10] }[variant]; // 2. Pre-allocate static tensors const tensors = [tensor('float32', outShape)] as const; diff --git a/.agents/skills/core-guidelines/SKILL.md b/.agents/skills/core-guidelines/SKILL.md index 16028711ff..2510c5b475 100644 --- a/.agents/skills/core-guidelines/SKILL.md +++ b/.agents/skills/core-guidelines/SKILL.md @@ -72,7 +72,7 @@ Use the following index to locate the specific procedural guides for your task: | **Add a new native operator or C++ binding** | [SKILL.md](../add-native-extension/SKILL.md) | Procedural guide to implementing C++ functions, exposing them via JSI, and writing TypeScript bridge wrappers. | | **Create a task pipeline or hook** | [SKILL.md](../add-task-pipeline/SKILL.md) | Guide to building end-to-end TS pipelines (e.g. object detection) and exposing them via React hooks. | | **Verify, rebuild, or troubleshoot changes** | [SKILL.md](../verify-and-build/SKILL.md) | Workflows for rebuilding TS/C++ and resolving common JSI runtime errors. | -| **Validate model constraints & schemas** | [SKILL.md](../model-schema-validation/SKILL.md) | Guide on specifying SymbolicTensor constraints and shapes for model signature validation. | +| **Validate model constraints & schemas** | [SKILL.md](../model-schema-validation/SKILL.md) | Guide on specifying model specs, dynamic shapes, and runtime constraints for model validation. | | **Maintain or refactor codebase patterns** | [SKILL.md](../skills-maintenance/SKILL.md) | Guide to keeping workspace skills in sync with codebase state to prevent documentation decay. | --- diff --git a/.agents/skills/model-schema-validation/SKILL.md b/.agents/skills/model-schema-validation/SKILL.md index f9d5e56ea0..c102fd72ef 100644 --- a/.agents/skills/model-schema-validation/SKILL.md +++ b/.agents/skills/model-schema-validation/SKILL.md @@ -1,162 +1,231 @@ --- name: model-schema-validation -description: Use when defining SymbolicTensor constraints, validating ExecuTorch model shapes, checking method signatures, or verifying dynamic tensor dimensions. +description: Use when defining model spec constraints, validating ExecuTorch model shapes, checking method signatures, or resolving dimension symbols. metadata: id: model_schema_validation - scope: src/core/modelSchema.ts, src/extensions/*/tasks/* + scope: src/core/schema.ts, src/extensions/*/tasks/* --- # Skill: Model Schema Validation Guide -Use this guide to define and enforce structural constraints (shapes and data types) on loaded ExecuTorch `.pte` models using `validateModelSchema`. +Use this guide to define and enforce structural constraints (shapes, data types, runtime constraints) on loaded ExecuTorch `.pte` models using `validateSpec`. --- -## 🔍 Why Validate Model Schemas? +## 🔍 Why Validate Model Specs? -Every ExecuTorch model exposes metadata about its execution methods (typically `'forward'`), including: -* Input and output argument counts. -* The expected types (primitives like `number`/`boolean` or `Tensor`). -* The data type (`float32`, `int32`, etc.) and the shape arrays of tensors. +Every ExecuTorch model exposes a `schema` (`ModelSpec`), derived either from standard ExecuTorch `MethodMeta` at load time or from an exported `get_model_schema` companion method returning a JSON model spec. -To prevent runtime crashes and memory allocation mismatches in C++, the TypeScript task pipeline must validate that the provided model matches its expected execution signature *before* allocating static tensors or executing inference. +To prevent runtime crashes, type mismatches, and memory allocation errors in C++, TypeScript task pipelines validate the loaded model's exported schema against allowed spec variants using `validateSpec` _before_ allocating static tensors or executing inference. --- ## 🛠️ Validation API Reference ```typescript -import { validateModelSchema, SymbolicTensor } from '../../../core/modelSchema'; - -const meta = validateModelSchema( - model, - methodName, - expectedInputs, - expectedOutputs -); +import { + validateSpec, + method, + f32, + i64, + i32, + DynamicDim as Dyn, + constr, +} from '../../../core/schema'; + +const { variant, dims } = validateSpec(model.schema, { + batched: method( + 'forward', + [i64(1, Dyn('L')), i64(1, Dyn('L'))], + [f32(1, 'D')], + [ + constr.eq( + { paramSide: 'input', tensorIdx: 0, dimIdx: 1 }, + { paramSide: 'input', tensorIdx: 1, dimIdx: 1 } + ), + ] + ), + unbatched: method( + 'forward', + [i64(Dyn('L')), i64(Dyn('L'))], + [f32('D')], + [ + constr.eq( + { paramSide: 'input', tensorIdx: 0, dimIdx: 0 }, + { paramSide: 'input', tensorIdx: 1, dimIdx: 0 } + ), + ] + ), +}); + +const [D] = dims.constant('D'); +const L = dims.range('L'); ``` -### Constraints Types: -* **Primitives**: `'number' | 'boolean' | 'null'` -* **Tensors**: Defined via the `SymbolicTensor(dtype?, ...shapes)` helper. +### Key Schema Utilities from `src/core/schema.ts`: + +- **`method(name, inputs, outputs, runtimeConstraints?)`**: Constructs a method specification. +- **`f32(...)` / `i64(...)` / `i32(...)` / `ui8(...)`**: Shorthand helpers for tensor parameter specs. +- **`StaticDim('symbol')` / String Literals**: Strings passed to shape helpers (e.g. `'H'`, `'W'`) automatically map to `StaticDim`, acting as **static dimension wildcards**. They bind strictly to `constant` positive integer dimensions in the exported spec. +- **`DynamicDim('symbol')` (or `Dyn('symbol')`)**: Creates a dynamic dimension symbol. Must be used when a dimension genuinely varies at runtime and binds to a `range` or `enum` domain in the exported spec. +- **Constraint Helpers (`constr`)**: + - **`DimRef` Object Literal (`{ paramSide: 'input' | 'output', tensorIdx, dimIdx }`)**: Explicit reference to a tensor's dimension. + - **`constr.eq(...dims)`**: Creates an equality constraint requiring the referenced dimensions to take the exact same value at runtime. + - **`constr.linear(lhs, rhs, a, b)`**: Creates a linear constraint `lhs = a * rhs + b`. +- **`validateSpec(exportedSchema, allowedVariants)`**: Compares the model's exported schema against named variants and returns `{ variant, dim, dims }`. +- **Symbol Accessors (`dims` & `dim`)**: + - `dims.constant('N', 'H')`: Extracts constant values for symbols as numbers. + - `dims.range('S')`: Extracts range domains `{ min, max, step }`. + - `dims.enum('E')`: Extracts enum choices `readonly number[]`. + - `dims.dynamic('L')`: Extracts dynamic `range` or `enum` as raw `ConcreteDim`. + - `dims.any('D')`: Extracts raw dimension value (`number`, `Range`, `readonly number[]`, or `ConcreteDim`). + - `dim('N')`: Extracts value for a single symbol. --- ## 📏 Symbolic Dimensions & Dynamic Shapes -Tensors often support dynamic dimensions (such as varying image sizes `'H'`, `'W'` or dynamic batch/prediction counts `'N'`). -The `SymbolicTensor` helper supports specifying **Symbolic Shapes** where integers are static matching constraints, and string names act as runtime variables. +Tensors can have static dimensions (integers), static symbol wildcards (strings), or dynamic symbols (`DynamicDim`). +The `validateSpec` utility matches allowed variants in order: -### How Symbolic Matching Works: 1. **Numbers (Static Match)**: - If a dimension is defined as a number, the loaded model's tensor dimension must match that exact integer. - * *Example*: `[1, 3, 'H', 'W']` requires the batch dimension to be exactly `1`, and channels to be exactly `3`. + Must match the exact static integer in the exported spec. -2. **Strings (Symbolic Match)**: - If a dimension is defined as a string (e.g., `'H'`), the validator binds the actual dimension size to that symbol name. - * *Symbolic Constraint Rules*: Within a single tensor, if a symbol (e.g., `'H'`) appears multiple times, the corresponding dimensions must be equal. +2. **Strings / `StaticDim` (Static Wildcard Match)**: + Binds string symbol names (e.g. `'H'`) to static constant integers in the exported spec. Repeated occurrences of the same static symbol must bind to the same constant value. -3. **Multiple Alternative Shapes**: - You can provide multiple shape variations to `SymbolicTensor` to support models compiled in different layouts (e.g., channels-first vs. channels-last). - * *Example*: `SymbolicTensor('float32', [1, 3, 'H', 'W'], [3, 'H', 'W'])` allows either 4D or 3D float32 layouts. +3. **`DynamicDim` (Dynamic Domain Match)**: + Binds dynamic symbol names to `range` or `enum` domains in the exported spec. + +4. **Variant Selection**: + Specify multiple variant keys in `validateSpec` (e.g. `batched` vs `unbatched`). The validator tests variants sequentially and returns the first matching key and bound symbols. --- -## ⚙️ Runtime Dynamic Dimensions: the `get_dynamic_dims_` companion +## 🔀 Dimension Domains vs. Runtime Constraints + +Understanding the distinction between a dimension's **domain** and its **runtime value** is central to schema design and validation: + +### 1. Dimension Domains & Domain Matching + +- **Dimension Domain**: The set of values a dimension may take: + - `constant`: Exactly `value` (a singleton integer). + - `range`: Any value between `min` and `max` stepping by `step`. + - `enum`: Any value listed in `choices`. +- **Runtime Value**: The concrete integer size of a tensor's dimension in a specific execution call. +- **Domain Matching (`validateSpec`)**: + - `validateSpec` performs static, load-time domain matching. + - Dynamic symbols (`DynamicDim('S')`) bind to exported dimension domains. Reusing a symbol (`'S'`) across tensor inputs or outputs requires every occurrence to bind to the **exact same domain**. + - ⚠️ **Key Rule**: Binding to the same domain does **NOT** mean runtime values coincide! Two dimensions bound to the same domain (e.g., both having range `1..512`) may take _different_ runtime values in a single execution (e.g. length 10 and length 25). + +### 2. Runtime Constraints (`constr.eq` & `constr.linear`) + +- **Runtime Constraints**: Declarations about the **runtime values** of tensor dimensions during execution: + - **Equality Constraint (`constr.eq(...)`)**: Requires all referenced dimensions to take the exact same runtime value in any execution call. + ```typescript + method( + 'forward', + [f32('B', Dyn('S1')), f32('B', Dyn('S2'))], + [f32('B', Dyn('S1'))], + [ + constr.eq( + { paramSide: 'input', tensorIdx: 0, dimIdx: 0 }, + { paramSide: 'input', tensorIdx: 1, dimIdx: 0 } + ), + ] + ); + ``` + - **Linear Constraint (`constr.linear(...)`)**: Requires two dimensions to satisfy `dimLhs = a * dimRhs + b` at runtime. +- **Validation & Enforcement**: + - `validateSpec` verifies that the exported model spec declares the exact same runtime constraints (1-to-1 declaration match). + - Native C++ validates input runtime constraints before invoking `model.execute()`. -`SymbolicTensor` string symbols let the **validator** accept a range of shapes, but an ExecuTorch `.pte`'s metadata only serializes the **static upper bound** of a dynamic dimension — not its active `[min, max, step]` range. So for an input dimension that genuinely varies at runtime (e.g. a text model's sequence length), the model must be **exported with a companion method** that re-exposes the range to the runtime: +--- -* **Name**: `get_dynamic_dims_` — e.g. `get_dynamic_dims_forward` for `forward`. -* **Signature**: takes no arguments. -* **Returns**: a list of outputs, one per `Tensor` input of the target method (scalar inputs are skipped), each a **2D `int32` tensor of shape `[rank, 3]`** whose rows are `[min, max, step]` constraints for that input's dimensions — e.g. `[1, 1, 1]` for a fixed dimension and `[1, maxTokens, 1]` for the dynamic one. +## ⚙️ Companion JSON Spec (`get_model_schema`) -At load time the C++ core (`Model::parseDynamicInputShapes`) reads this companion and validates inputs against the range; `model.execute` then accepts tensors at their exact runtime length. **Without the companion, a method falls back to exact per-dimension validation** — it only accepts the single shape it was exported with. So a `.pte` whose metadata reports `[1, 512]` but is meant to accept `[1, 1..512]` MUST ship `get_dynamic_dims_forward`, or variable-length inputs are rejected at runtime. +For models with runtime dynamic dimensions or runtime constraints (equality/linear relations between dimensions), the `.pte` model exports a companion method named `get_model_schema` returning a JSON string representation of `ModelSpec`. -This is an **export-side contract** (the export-scripts repo provides a `build_dynamic_dims_program` helper that emits the companion). The TypeScript task only declares the symbol via `SymbolicTensor` and reads the resulting upper bound from `meta.inputTensorMeta`. +When `get_model_schema` is present, the C++ loader parses it to populate `model.schema` with exact `RangeDim` / `EnumDim` domains and runtime constraints. Without it, `model.schema` is populated strictly from static ExecuTorch `MethodMeta`. --- ## 📋 Common Validation Recipes -### 1. Classification -Accepts an image tensor and outputs a 2D class probabilities array: +### 1. Classification (Batched vs Unbatched) + +Accepts an image tensor and outputs class probabilities: + ```typescript -const meta = validateModelSchema( - model, - 'forward', - [SymbolicTensor('float32', [1, 3, 'H', 'W'], [3, 'H', 'W'])], // Input - [SymbolicTensor('float32', [1, 'N'], ['N'])] // Output: logits / probs -); - -const inpShape = meta.inputTensorMeta[0]!.shape; -const outShape = meta.outputTensorMeta[0]!.shape; +const { variant, dims } = validateSpec(model.schema, { + batched: method('forward', [f32(1, 3, 'H', 'W')], [f32(1, 'N')]), + unbatched: method('forward', [f32(3, 'H', 'W')], [f32('N')]), +}); + +const [N, H, W] = dims.constant('N', 'H', 'W'); +const inpShape = { batched: [1, 3, H, W], unbatched: [3, H, W] }[variant]; +const outShape = { batched: [1, N], unbatched: [N] }[variant]; ``` ### 2. Image-to-Image / Style Transfer + Accepts an image tensor and returns a modified image tensor with identical dimensions: + ```typescript -const meta = validateModelSchema( - model, - 'forward', - [SymbolicTensor('float32', [1, 3, 'H', 'W'], [3, 'H', 'W'])], // Input - [SymbolicTensor('float32', [1, 3, 'H', 'W'], [3, 'H', 'W'])] // Output -); +const { dims } = validateSpec(model.schema, { + default: method('forward', [f32(1, 3, 'H', 'W')], [f32(1, 3, 'H', 'W')]), +}); + +const [H, W] = dims.constant('H', 'W'); ``` -### 3. Object Detection (Dynamic Boxes Count) -Accepts an image tensor, and outputs boxes, scores, and class labels for `N` dynamic detections: +### 3. Object Detection (Dynamic Box Count) + +Accepts an image tensor, and outputs boxes, scores, and class labels for `N` detections: + ```typescript -const meta = validateModelSchema( - model, - 'forward', - [SymbolicTensor('float32', [1, 3, 'H', 'W'], [3, 'H', 'W'])], - [ - SymbolicTensor('float32', ['N', 4]), // Bounding boxes (xyxy / xywh) - SymbolicTensor('float32', ['N']), // Prediction confidence scores - SymbolicTensor('float32', ['N']), // Predicted class labels - ] -); +const { dims } = validateSpec(model.schema, { + default: method('forward', [f32(1, 3, 'H', 'W')], [f32('N', 4), f32('N'), f32('N')]), +}); + +const [N, H, W] = dims.constant('N', 'H', 'W'); ``` -### 4. Single-Model Pipeline (fully static export) +### 4. Single-Model Pipeline (Fully Static Export) -When a pipeline targets exactly one model whose `.pte` is exported with fully static shapes (e.g. `sdxsTextToImage`), the shapes are constants of the export contract. Validate each method against those constants — do **not** invent symbols for dimensions that cannot vary, and do **not** thread the shapes through task options: +Assert exact shape constants from the task file: ```typescript -const IMAGE_SIZE = 512; -const LATENT_CHANNELS = 4; -const LATENT_SIZE = IMAGE_SIZE / 8; - -validateModelSchema( - model, - 'denoise', - [ - SymbolicTensor('float32', [1, LATENT_CHANNELS, LATENT_SIZE, LATENT_SIZE]), - SymbolicTensor('int64', [1]), - SymbolicTensor('float32', [1, CLIP_MAX_TOKENS, CLIP_HIDDEN_SIZE]), - ], - [SymbolicTensor('float32', [1, LATENT_CHANNELS, LATENT_SIZE, LATENT_SIZE])] -); +const { dims } = validateSpec(model.schema, { + default: method( + 'denoise', + [ + f32(1, LATENT_CHANNELS, LATENT_SIZE, LATENT_SIZE), + i64(1), + f32(1, CLIP_MAX_TOKENS, CLIP_HIDDEN_SIZE), + ], + [f32(1, LATENT_CHANNELS, LATENT_SIZE, LATENT_SIZE)] + ), +}); ``` -A rigid `.pte` contract is an argument **for** validating, not against it: the stricter the contract, the cheaper and more precise the assertion. Validation is load-time only (it costs nothing per inference) and converts a stale, corrupt, or wrongly re-exported `.pte` into a readable TypeScript error instead of a native crash or silently garbage output. This matters most for pipelines shipping several separately-exported backend variants (XNNPACK / CoreML / MLX), where one variant can break while the others stay healthy. - --- ## 🚫 Avoid / Anti-Patterns -* **Do NOT write imperative size/type checks manually:** Avoid writing custom shape/type assertion blocks (e.g., `if (tensor.shape[0] !== 1)`). Always use the declarative `validateModelSchema` utility, which reports unified, readable mismatch errors. -* **Do NOT use hardcoded integers for dynamic dimensions:** If a shape can vary (e.g., dynamic height, width, or batch sizes), use a string symbol (like `'H'`, `'W'`, `'N'`) to allow dynamic matching. Conversely, a dimension that genuinely cannot vary *should* be a hardcoded integer — reserve symbols for real variability. -* **Do NOT skip validation at startup:** Always validate the model schema *before* creating pre-allocated static tensors, preventing native memory crashes from mismatched layouts. -* **Do NOT skip validation just because the pipeline supports a single model:** "The `.pte` must be very specific anyway" is a reason to assert the exact contract, not to trust it. See Recipe 4. +- **Do NOT write manual imperative shape/type checks:** Always use declarative `validateSpec` variants which report unified, readable mismatch errors. +- **Do NOT use hardcoded integers for dynamic dimensions:** Use string symbols (like `'H'`, `'W'`, `'N'`) for dynamic dimensions. Conversely, fixed export dimensions should be integers. +- **Do NOT skip validation at startup:** Validate `model.schema` before allocating static tensors. +- **Do NOT skip validation for single-model pipelines:** Single-model pipelines should still validate against task shape constants. --- ## 📋 Verification Checklist When specifying model schema validations, verify that: -- [ ] Schema validation is performed immediately after model loading and before tensor initialization. -- [ ] All dynamic dimensions (e.g., dynamic box counts, channels-last heights/widths) are defined as string symbols, and dimensions fixed by the export are plain integers. -- [ ] Single-model pipelines with static exports still validate, asserting against the task file's shape constants rather than skipping validation or routing shapes through options. -- [ ] Multiple shape variations are provided to `SymbolicTensor` if channels-first and channels-last layouts are both supported. -- [ ] Input and output constraints map accurately to standard model specifications (e.g. dense logits, standard bounding boxes layouts). + +- [ ] Schema validation is performed immediately after model loading and before tensor initialization using `validateSpec(model.schema, { ... })`. +- [ ] Dynamic dimensions are defined as string symbols, while fixed dimensions are plain integers. +- [ ] Symbol values are extracted using `dims.constant(...)`, `dims.range(...)`, or `dims.enum(...)`. +- [ ] Multiple shape variants (e.g. `batched` vs `unbatched`) are provided when supported. +- [ ] Input and output constraints map accurately to model specifications. diff --git a/apps/computer-vision/app/inspect/index.tsx b/apps/computer-vision/app/inspect/index.tsx index 5085ec0d31..b1505028bd 100644 --- a/apps/computer-vision/app/inspect/index.tsx +++ b/apps/computer-vision/app/inspect/index.tsx @@ -9,12 +9,25 @@ import { ActivityIndicator, Alert, } from 'react-native'; -import { inspectModel, type TensorMeta } from 'react-native-executorch'; +import { inspectModel, type ConcreteDim, type ParamSpec } from 'react-native-executorch'; import ScreenWrapper from '../../components/ScreenWrapper'; import { ColorPalette } from '../../theme'; type InspectionResult = Awaited>; +const formatDim = (dim: ConcreteDim): string => { + switch (dim.kind) { + case 'constant': + return `${dim.value}`; + case 'range': + return dim.range.step !== 1 + ? `${dim.range.min}..${dim.range.max} (step ${dim.range.step})` + : `${dim.range.min}..${dim.range.max}`; + case 'enum': + return dim.choices.join(' | '); + } +}; + function InspectContent() { const [url, setUrl] = useState(''); const [loading, setLoading] = useState(false); @@ -39,27 +52,32 @@ function InspectContent() { } }; - const renderTensorList = (tensors: TensorMeta[] | undefined, title: string) => { - if (!tensors || tensors.length === 0) return null; + const renderParamList = ( + params: readonly ParamSpec[] | undefined, + title: string + ) => { + if (!params || params.length === 0) return null; return ( {title} - {tensors.map((tensor, idx) => ( + {params.map((param, idx) => ( - - {tensor.name || `${title.slice(0, -1)} #${idx}`} - - {tensor.dtype} - - - - Shape: [{tensor.shape.join(', ')}] - - - Bytes: {tensor.nbytes} + #{idx} + + {param.kind === 'Tensor' ? param.dtype : param.kind} + {param.kind === 'Tensor' && ( + + + Shape:{' '} + + [{param.shape.map(formatDim).join(', ')}] + + + + )} ))} @@ -113,12 +131,12 @@ function InspectContent() { Source URL: {result.source} - Methods ({result.methods.length}) + Methods ({Object.keys(result.schema).length}) - {result.methods.map((method, mIdx) => ( - + {Object.entries(result.schema).map(([methodName, spec]) => ( + - {method.name} + {methodName} Method @@ -126,35 +144,30 @@ function InspectContent() { - {method.meta.numInputs} + {spec.inputs.length} Inputs - {method.meta.numOutputs} + {spec.outputs.length} Outputs - {method.meta.usesBackend && Object.keys(method.meta.usesBackend).length > 0 && ( + {(result.backends[methodName]?.length ?? 0) > 0 && ( Backends Used: - {Object.entries(method.meta.usesBackend).map(([backend, used]) => ( - - - {backend}: {used ? 'Yes' : 'No'} - + {result.backends[methodName]?.map((backend) => ( + + {backend} ))} )} - {renderTensorList(method.meta.inputTensorMeta, 'Input Tensors')} - {renderTensorList(method.meta.outputTensorMeta, 'Output Tensors')} + {renderParamList(spec.inputs, 'Inputs')} + {renderParamList(spec.outputs, 'Outputs')} ))} diff --git a/packages/react-native-executorch/cpp/core/model.cpp b/packages/react-native-executorch/cpp/core/model.cpp index 5b9503cb82..2ff28cd2c3 100644 --- a/packages/react-native-executorch/cpp/core/model.cpp +++ b/packages/react-native-executorch/cpp/core/model.cpp @@ -10,6 +10,7 @@ #include #include "dtype.h" +#include "schema.h" #include "tensor_helpers.h" #include @@ -20,8 +21,6 @@ namespace { namespace jsi = facebook::jsi; -namespace types = rnexecutorch::core::types; -namespace tensor = rnexecutorch::core::tensor; template T unwrap(const std::string &ctx, executorch::runtime::Result result) { @@ -39,155 +38,6 @@ T unwrap(jsi::Runtime &rt, const std::string &ctx, executorch::runtime::Result(tensorMeta.sizes().size())); - jsTensorMeta.setProperty(rt, "nbytes", static_cast(tensorMeta.nbytes())); - - try { - auto dtypeStr = types::toString(types::fromScalarType(tensorMeta.scalar_type())); - jsTensorMeta.setProperty(rt, "dtype", jsi::String::createFromUtf8(rt, dtypeStr)); - } catch (const std::exception &) { - jsTensorMeta.setProperty(rt, "dtype", jsi::String::createFromUtf8(rt, "not supported")); - } - - auto jsShapeArray = jsi::Array(rt, tensorMeta.sizes().size()); - for (size_t i = 0; i < tensorMeta.sizes().size(); ++i) { - jsShapeArray.setValueAtIndex(rt, i, static_cast(tensorMeta.sizes()[i])); - } - jsTensorMeta.setProperty(rt, "shape", jsShapeArray); - - return jsTensorMeta; -} - -/** - * Parses compile-time dynamic dimension constraints for the inputs of a module - * method. - * - * @note This is a temporary workaround. ExecuTorch (.pte) model metadata - * natively serializes only the static upper-bound limits of dynamic/symbolic - * dimensions, and does not expose the active dynamic range (min, max, step) at - * runtime. - * - * To use this feature, the .pte model must be exported with a companion method - * named "get_dynamic_dims_" (e.g., "get_dynamic_dims_forward"). - * - * Python export requirements: - * 1. The companion method must take no arguments. - * 2. It must return a list of outputs matching the number of Tag::Tensor inputs - * of the target method (scalar inputs are excluded). - * 3. Each output must be a 2D int32 tensor of shape [rank, 3], where each row - * represents [min, max, step] constraints for the corresponding dimension of - * that input tensor. - * - * @param module The ExecuTorch extension module to query. - * @param methodName The name of the target module method (e.g. "forward"). - * @return A vector of SymbolicShape objects, or std::nullopt if the companion - * method is not defined. If returned, the vector's size is guaranteed - * to equal methodMeta.num_inputs(), containing parsed SymbolicShapes - * for tensor inputs and empty SymbolicShapes for non-tensor inputs. - */ -std::optional> -parseDynamicInputShapes(executorch::extension::Module &module, const std::string &methodName) { - using executorch::aten::ScalarType; - - const auto getDynamicShapesMethodName = std::format("get_dynamic_dims_{}", methodName); - const auto ctx = getDynamicShapesMethodName + ": "; - - auto methodNames = unwrap(ctx + "failed to get method names", module.method_names()); - if (methodName == getDynamicShapesMethodName || - !methodNames.contains(getDynamicShapesMethodName)) { - return std::nullopt; - } - - auto methodMeta = unwrap(std::format("{}failed to get meta for method '{}'", ctx, methodName), - module.method_meta(methodName)); - - size_t expectedTensorInputs = 0; - for (size_t i = 0; i < methodMeta.num_inputs(); ++i) { - auto tag = unwrap(std::format("{}failed to get tag for input [{}]", ctx, i), methodMeta.input_tag(i)); - if (tag == executorch::runtime::Tag::Tensor) { - expectedTensorInputs++; - } - } - - auto result = unwrap(ctx + "failed to execute", module.execute(getDynamicShapesMethodName)); - if (result.size() != expectedTensorInputs) { - throw std::runtime_error(std::format("{}number of outputs returned ({}) does not match the number of " - "tensor inputs declared by method '{}' ({})", - ctx, result.size(), methodName, expectedTensorInputs)); - } - - std::vector dynamicShapes; - dynamicShapes.reserve(methodMeta.num_inputs()); - size_t tensorIndex = 0; - - for (size_t i = 0; i < methodMeta.num_inputs(); ++i) { - auto tag = unwrap(std::format("{}failed to get tag for input [{}]", ctx, i), - methodMeta.input_tag(i)); - - if (tag != executorch::runtime::Tag::Tensor) { - // Emplace an empty shape for non-tensor inputs to maintain a 1-to-1 - // index alignment between dynamicShapes and the model's inputs vector. - dynamicShapes.emplace_back(); - continue; - } - - const auto &out = result.at(tensorIndex); - if (!out.isTensor()) { - throw std::runtime_error(std::format("{}output[{}] is not a tensor", ctx, tensorIndex)); - } - - auto inputMeta = unwrap(std::format("{}failed to get tensor meta for input [{}]", ctx, i), - methodMeta.input_tensor_meta(i)); - const auto rank = inputMeta.sizes().size(); - const auto shapeTensor = out.toTensor(); - - if (shapeTensor.dim() != 2 || shapeTensor.size(1) != 3 || - shapeTensor.size(0) != static_cast(rank) || - shapeTensor.scalar_type() != ScalarType::Int) { - throw std::runtime_error(std::format("{}output[{}] expected to be a 2D int32_t tensor of shape [{}, 3]", - ctx, tensorIndex, rank)); - } - - const auto *shape = shapeTensor.const_data_ptr(); - tensor::SymbolicShape symbolicShape; - symbolicShape.reserve(rank); - - for (size_t axis = 0; axis < rank; ++axis) { - const auto minDim = shape[axis * 3 + 0]; - const auto maxDim = shape[axis * 3 + 1]; - const auto step = shape[axis * 3 + 2]; - if (minDim < 0 || maxDim < minDim || step < 1) { - throw std::runtime_error(std::format("{}output[{}], axis {} is invalid: " - "expected 0 <= min <= max and step >= 1 but got [{}, {}, {}]", - ctx, tensorIndex, axis, minDim, maxDim, step)); - } - if (maxDim > inputMeta.sizes()[axis]) { - throw std::runtime_error(std::format("{}output[{}], axis {} max dimension ({}) " - "exceeds model metadata upper limit ({})", - ctx, tensorIndex, axis, maxDim, inputMeta.sizes()[axis])); - } - - symbolicShape.emplace_back(tensor::RangeDim{.min = minDim, .max = maxDim, .step = step}); - } - - dynamicShapes.push_back(std::move(symbolicShape)); - ++tensorIndex; - } - - return dynamicShapes; -} } // namespace namespace rnexecutorch::core::model { @@ -199,18 +49,41 @@ using rnexecutorch::core::tensor::TensorHostObject; ModelHostObject::ModelHostObject(const std::string &modelPath) : modelPath_(modelPath), etModule_(std::make_unique(modelPath)) { + auto error = etModule_->load(); if (!etModule_->is_loaded()) { const std::string errorMsg = executorch::runtime::to_string(error); - throw std::runtime_error(std::format("Failed to load model: {}", errorMsg)); + throw std::runtime_error(std::format("Failed to load model from '{}': {}", + modelPath_, errorMsg)); + } + + const auto methodNames = unwrap("method names", etModule_->method_names()); + schema::ModelSpec overrideSpec; + + if (methodNames.contains(kGetModelSchemaMethod)) { + auto ctx = std::format("Execute '{}'", kGetModelSchemaMethod); + auto result = unwrap(ctx, etModule_->execute(kGetModelSchemaMethod)); + + if (result.empty() || result[0].tag != executorch::runtime::Tag::String) { + throw std::runtime_error(std::format("{} must return a single string value", ctx)); + } + + auto jsonStr = std::string(result[0].toString()); + overrideSpec = schema::parseModelSpecJson(ctx, jsonStr); } - auto methodNames = unwrap("Failed to get method names", etModule_->method_names()); for (const auto &methodName : methodNames) { - auto dynamicShapes = parseDynamicInputShapes(*etModule_, methodName); - if (dynamicShapes) { - dynamicInputShapes_.emplace(methodName, std::move(*dynamicShapes)); + auto ctx = std::format("Method '{}'", methodName); + auto methodMeta = unwrap(ctx, etModule_->method_meta(methodName)); + + spec_[methodName] = schema::methodSpecFromMetadata(methodMeta); + backends_[methodName] = schema::getUsedBackends(methodMeta); + + if (overrideSpec.contains(methodName)) { + spec_[methodName] = std::move(overrideSpec[methodName]); } + + schema::validateSpec(spec_[methodName], methodMeta, ctx + " metadata validation"); } } @@ -221,105 +94,12 @@ jsi::Value ModelHostObject::get(jsi::Runtime &rt, const jsi::PropNameID &name) { return jsi::String::createFromUtf8(rt, modelPath_); } - if (nameStr == "getMethodNames") { - auto self = shared_from_this(); - auto fnBody = [self](jsi::Runtime &rt, const jsi::Value & /*thisVal*/, const jsi::Value * /*args*/, size_t count) -> jsi::Value { - if (count != 0) { - throw jsi::JSError(rt, "getMethodNames: Usage: getMethodNames()"); - } - - std::unique_lock lock(self->mutex_, std::try_to_lock); - if (!lock.owns_lock()) { - throw jsi::JSError(rt, "getMethodNames: Model is currently in use"); - } - - if (!self->etModule_) { - throw jsi::JSError(rt, "getMethodNames: Model has been disposed"); - } - - auto methodNames = unwrap(rt, "getMethodNames", self->etModule_->method_names()); - - auto jsArray = jsi::Array(rt, methodNames.size()); - size_t index = 0; - for (const auto &methodName : methodNames) { - jsArray.setValueAtIndex(rt, index, jsi::String::createFromUtf8(rt, methodName)); - ++index; - } - - return jsArray; - }; - return jsi::Function::createFromHostFunction(rt, jsi::PropNameID::forAscii(rt, "getMethodNames"), 0, fnBody); + if (nameStr == "schema") { + return schema::modelSpecToJs(rt, spec_); } - if (nameStr == "getMethodMeta") { - auto self = shared_from_this(); - auto fnBody = [self](jsi::Runtime &rt, const jsi::Value & /*thisVal*/, const jsi::Value *args, size_t count) -> jsi::Value { - using executorch::runtime::tag_to_string; - - if (count != 1) { - throw jsi::JSError(rt, "getMethodMeta: Usage: getMethodMeta(methodName)"); - } - - std::unique_lock lock(self->mutex_, std::try_to_lock); - if (!lock.owns_lock()) { - throw jsi::JSError(rt, "getMethodMeta: Model is currently in use"); - } - - if (!self->etModule_) { - throw jsi::JSError(rt, "getMethodMeta: Model has been disposed"); - } - - auto methodName = conversions::asType(rt, "getMethodMeta: methodName", args[0]); - auto methodMeta = unwrap(rt, "getMethodMeta", self->etModule_->method_meta(methodName)); - - auto inputTagsArray = jsi::Array(rt, methodMeta.num_inputs()); - for (size_t i = 0; i < methodMeta.num_inputs(); ++i) { - auto ctx = std::format("getMethodMeta: input tag [{}]", i); - auto tag = unwrap(rt, ctx, methodMeta.input_tag(i)); - inputTagsArray.setValueAtIndex(rt, i, jsi::String::createFromUtf8(rt, tag_to_string(tag))); - } - - auto outputTagsArray = jsi::Array(rt, methodMeta.num_outputs()); - for (size_t i = 0; i < methodMeta.num_outputs(); ++i) { - auto ctx = std::format("getMethodMeta: output tag [{}]", i); - auto tag = unwrap(rt, ctx, methodMeta.output_tag(i)); - outputTagsArray.setValueAtIndex(rt, i, jsi::String::createFromUtf8(rt, tag_to_string(tag))); - } - - auto usesBackendMap = jsi::Object(rt); - for (size_t i = 0; i < methodMeta.num_backends(); ++i) { - auto ctx = std::format("getMethodMeta: backend name [{}]", i); - const auto *backendName = unwrap(rt, ctx, methodMeta.get_backend_name(i)); - usesBackendMap.setProperty(rt, backendName, methodMeta.uses_backend(backendName)); - } - - auto inputTensorMetaArray = jsi::Array(rt, methodMeta.num_inputs()); - for (size_t i = 0; i < methodMeta.num_inputs(); ++i) { - auto ctx = std::format("getMethodMeta: input tensor meta [{}]", i); - auto tensorMeta = unwrap(rt, ctx, methodMeta.input_tensor_meta(i)); - inputTensorMetaArray.setValueAtIndex(rt, i, tensorMetaToJs(rt, tensorMeta)); - } - - auto outputTensorMetaArray = jsi::Array(rt, methodMeta.num_outputs()); - for (size_t i = 0; i < methodMeta.num_outputs(); ++i) { - auto ctx = std::format("getMethodMeta: output tensor meta [{}]", i); - auto tensorMeta = unwrap(rt, ctx, methodMeta.output_tensor_meta(i)); - outputTensorMetaArray.setValueAtIndex(rt, i, tensorMetaToJs(rt, tensorMeta)); - } - - auto jsMeta = jsi::Object(rt); - jsMeta.setProperty(rt, "name", jsi::String::createFromUtf8(rt, methodMeta.name())); - jsMeta.setProperty(rt, "numInputs", static_cast(methodMeta.num_inputs())); - jsMeta.setProperty(rt, "numOutputs", static_cast(methodMeta.num_outputs())); - jsMeta.setProperty(rt, "inputTags", inputTagsArray); - jsMeta.setProperty(rt, "outputTags", outputTagsArray); - jsMeta.setProperty(rt, "usesBackend", usesBackendMap); - jsMeta.setProperty(rt, "inputTensorMeta", inputTensorMetaArray); - jsMeta.setProperty(rt, "outputTensorMeta", outputTensorMetaArray); - - return jsMeta; - }; - return jsi::Function::createFromHostFunction(rt, jsi::PropNameID::forAscii(rt, "getMethodMeta"), 1, fnBody); + if (nameStr == "backends") { + return schema::backendsToJs(rt, backends_); } if (nameStr == "execute") { @@ -339,43 +119,40 @@ jsi::Value ModelHostObject::get(jsi::Runtime &rt, const jsi::PropNameID &name) { } auto methodName = conversions::asType(rt, "execute: methodName", args[0]); - auto methodMeta = unwrap(rt, "execute", self->etModule_->method_meta(methodName)); + if (!self->spec_.contains(methodName)) { + throw jsi::JSError(rt, std::format("execute: Unknown method '{}'", methodName)); + } + const auto &methodSpec = self->spec_.at(methodName); auto inputsArray = conversions::asType(rt, "execute: inputs", args[1]); auto outputTensorsArray = conversions::asType(rt, "execute: outputTensors", args[2]); - if (inputsArray.size(rt) != methodMeta.num_inputs()) { - throw jsi::JSError(rt, std::format("execute: Incorrect size for inputs: got {}, expected {}", - inputsArray.size(rt), methodMeta.num_inputs())); + if (inputsArray.size(rt) != methodSpec.inputs.size()) { + throw jsi::JSError(rt, std::format("execute: Incorrect size for inputs of method '{}': got {}, expected {}", + methodName, inputsArray.size(rt), methodSpec.inputs.size())); } - std::vector inputs(methodMeta.num_inputs()); + std::vector inputs(methodSpec.inputs.size()); std::vector> tensorLocks; std::unordered_set lockedTensors; + std::vector> inputShapes; - for (size_t i = 0; i < methodMeta.num_inputs(); ++i) { + for (size_t i = 0; i < methodSpec.inputs.size(); ++i) { auto ctx = std::format("execute: inputs[{}]", i); - auto tag = unwrap(rt, ctx, methodMeta.input_tag(i)); + auto tag = methodSpec.inputs[i].tag; auto val = inputsArray.getValueAtIndex(rt, i); switch (tag) { case executorch::runtime::Tag::Tensor: { - auto tensorMeta = unwrap(rt, ctx + ": tensor meta", methodMeta.input_tensor_meta(i)); - auto expectedDtype = fromScalarType(rt, ctx, tensorMeta.scalar_type()); - - std::shared_ptr tensorHostObject; - if (self->dynamicInputShapes_.contains(methodName)) { - auto expectedShape = self->dynamicInputShapes_[methodName][i]; - tensorHostObject = tensor::fromJs(rt, ctx, val, expectedDtype, expectedShape); - } else { - tensorHostObject = tensor::fromJs(rt, ctx, val, expectedDtype, tensorMeta.sizes()); - } + const auto &tSpec = methodSpec.inputs[i]; + auto tensorHostObject = tensor::fromJs(rt, ctx, val, tSpec.dtype, tSpec.shape); if (!lockedTensors.insert(tensorHostObject.get()).second) { throw jsi::JSError(rt, "execute: Tensor aliasing detected. " "The same tensor was passed multiple times."); } tensorLocks.emplace_back(tensor::tryLockUnique(rt, ctx, tensorHostObject)); + inputShapes.push_back(tensorHostObject->shape_); inputs[i] = tensorHostObject->tensor_; break; } @@ -397,6 +174,9 @@ jsi::Value ModelHostObject::get(jsi::Runtime &rt, const jsi::PropNameID &name) { } } + schema::validateRuntimeConstraints(rt, methodSpec.runtimeConstraints, inputShapes, + std::format("execute: method '{}'", methodName)); + auto startTime = std::chrono::high_resolution_clock::now(); auto executeResult = self->etModule_->execute(methodName, inputs); auto finishTime = std::chrono::high_resolution_clock::now(); @@ -409,30 +189,49 @@ jsi::Value ModelHostObject::get(jsi::Runtime &rt, const jsi::PropNameID &name) { logFn.callWithThis(rt, consoleObj, {jsi::String::createFromUtf8(rt, info)}); #endif - auto result = unwrap(rt, - std::format("execute: Method '{}' failed (check getMethodMeta() " - "for required backends and getRegisteredBackends() " - "for registered ones)", - methodName), + auto result = unwrap(rt, std::format("execute: Method '{}' failed.\n" + "\n" + "Common causes:\n" + " 1. Backend not registered\n" + " Ensure backends from `model.backends` are registered\n" + " in the ExecuTorch runtime\n" + " (use `getRegisteredBackends()` to check registered backends).\n" + "\n" + " 2. Shape/constraint mismatch\n" + " If the model uses dynamic shapes or runtime constraints\n" + " (e.g. equality between dimensions), export a companion\n" + " method returning a JSON model spec\n" + " (see `src/core/schema.ts` for the JSON structure).\n" + " Without it, validation falls back to static metadata\n" + " from ExecuTorch which only contains upper bounds and\n" + " does not capture runtime constraints.\n" + "\n" + " 3. Bad model export\n" + " The model export itself might be broken or invalid.\n" + "\n" + "Error", + methodName), std::move(executeResult)); auto jsOutputArray = jsi::Array(rt, result.size()); - size_t index = 0; + size_t outputIdx = 0; size_t tensorOutputIdx = 0; for (const auto &output : result) { switch (output.tag) { case executorch::runtime::Tag::Tensor: { if (tensorOutputIdx >= outputTensorsArray.size(rt)) { - throw jsi::JSError(rt, "execute: Not enough tensor output placeholders in outputTensors"); + throw jsi::JSError(rt, std::format("execute: Not enough tensor output placeholders in outputTensors" + " (provided {}, expected at least {})", + outputTensorsArray.size(rt), tensorOutputIdx + 1)); } auto ctx = std::format("execute: outputTensors[{}]", tensorOutputIdx); auto val = outputTensorsArray.getValueAtIndex(rt, tensorOutputIdx); - auto tensorMeta = unwrap(rt, ctx + ": tensor meta", methodMeta.output_tensor_meta(index)); - auto expectedDtype = fromScalarType(rt, ctx, tensorMeta.scalar_type()); - auto tensorHostObject = tensor::fromJs(rt, ctx, val, expectedDtype, output.toTensor().sizes()); + auto dtype = types::fromScalarType(output.toTensor().dtype()); + auto shape = output.toTensor().sizes(); + auto tensorHostObject = tensor::fromJs(rt, ctx, val, dtype, shape); if (!lockedTensors.insert(tensorHostObject.get()).second) { throw jsi::JSError(rt, "execute: Tensor aliasing detected. " @@ -443,28 +242,31 @@ jsi::Value ModelHostObject::get(jsi::Runtime &rt, const jsi::PropNameID &name) { output.toTensor().const_data_ptr(), output.toTensor().nbytes()); - jsOutputArray.setValueAtIndex(rt, index, val); + jsOutputArray.setValueAtIndex(rt, outputIdx, val); ++tensorOutputIdx; break; } case executorch::runtime::Tag::Double: - jsOutputArray.setValueAtIndex(rt, index, output.toDouble()); + jsOutputArray.setValueAtIndex(rt, outputIdx, output.toDouble()); break; case executorch::runtime::Tag::Int: - jsOutputArray.setValueAtIndex(rt, index, static_cast(output.toInt())); + jsOutputArray.setValueAtIndex(rt, outputIdx, static_cast(output.toInt())); break; case executorch::runtime::Tag::Bool: - jsOutputArray.setValueAtIndex(rt, index, output.toBool()); + jsOutputArray.setValueAtIndex(rt, outputIdx, output.toBool()); break; case executorch::runtime::Tag::None: - jsOutputArray.setValueAtIndex(rt, index, jsi::Value::null()); + jsOutputArray.setValueAtIndex(rt, outputIdx, jsi::Value::null()); + break; + case executorch::runtime::Tag::String: + jsOutputArray.setValueAtIndex(rt, outputIdx, jsi::String::createFromUtf8(rt, std::string(output.toString()))); break; default: throw jsi::JSError(rt, std::format("execute: Unsupported return type: {}", executorch::runtime::tag_to_string(output.tag))); } - ++index; + ++outputIdx; } return jsOutputArray; @@ -498,8 +300,8 @@ jsi::Value ModelHostObject::get(jsi::Runtime &rt, const jsi::PropNameID &name) { std::vector ModelHostObject::getPropertyNames(jsi::Runtime &rt) { std::vector properties; properties.push_back(jsi::PropNameID::forAscii(rt, "path")); - properties.push_back(jsi::PropNameID::forAscii(rt, "getMethodNames")); - properties.push_back(jsi::PropNameID::forAscii(rt, "getMethodMeta")); + properties.push_back(jsi::PropNameID::forAscii(rt, "schema")); + properties.push_back(jsi::PropNameID::forAscii(rt, "backends")); properties.push_back(jsi::PropNameID::forAscii(rt, "execute")); properties.push_back(jsi::PropNameID::forAscii(rt, "dispose")); return properties; diff --git a/packages/react-native-executorch/cpp/core/model.h b/packages/react-native-executorch/cpp/core/model.h index b196d48664..594e28679b 100644 --- a/packages/react-native-executorch/cpp/core/model.h +++ b/packages/react-native-executorch/cpp/core/model.h @@ -6,6 +6,7 @@ #include #include +#include "schema.h" #include "tensor_helpers.h" #include @@ -14,6 +15,13 @@ namespace rnexecutorch::core::model { namespace jsi = facebook::jsi; + +/** + * Optional companion method name inside a `.pte` ExecuTorch module that exports + * a JSON model spec (containing dynamic dimension domains or runtime constraints). + */ +inline constexpr auto kGetModelSchemaMethod = "get_model_schema"; + /** * JSI HostObject wrapping an ExecuTorch Model instance * (`executorch::extension::Module`). @@ -25,18 +33,43 @@ namespace jsi = facebook::jsi; class ModelHostObject : public jsi::HostObject, public std::enable_shared_from_this { public: + /** + * Loads the ExecuTorch model from the specified file path, initializes its + * method metadata, parses schemas, and populates backend delegates. + * + * @param modelPath Absolute file system path to the `.pte` model binary. + */ explicit ModelHostObject(const std::string &modelPath); jsi::Value get(jsi::Runtime &rt, const jsi::PropNameID &name) override; std::vector getPropertyNames(jsi::Runtime &rt) override; private: + /// File system path to the loaded ExecuTorch `.pte` model binary. std::string modelPath_; - std::unique_ptr etModule_; + + /** + * Parsed model schema mapping method names to their input/output parameter + * specs and runtime constraints. + */ + schema::ModelSpec spec_; + + /** + * Map of method names to the list of backend delegate identifiers used by + * each method. + */ + std::unordered_map> backends_; + + /** + * Mutex serializing access to `etModule_` to prevent concurrent inference + * execution on the same model. + */ std::mutex mutex_; - std::unordered_map> dynamicInputShapes_; + /// Unique pointer to the underlying ExecuTorch module instance. + std::unique_ptr etModule_; }; void install_loadModel(jsi::Runtime &rt, jsi::Object &module); + } // namespace rnexecutorch::core::model diff --git a/packages/react-native-executorch/cpp/core/schema.cpp b/packages/react-native-executorch/cpp/core/schema.cpp new file mode 100644 index 0000000000..e26e6dfd5e --- /dev/null +++ b/packages/react-native-executorch/cpp/core/schema.cpp @@ -0,0 +1,535 @@ +#include "schema.h" + +#include +#include +#include + +#include +#include + +#include "dtype.h" + +namespace nlohmann { +template <> +// Tag lives in executorch::runtime; adl_serializer is the +// nlohmann-blessed extension point for types we don't own. +// See: https://github.com/nlohmann/json#how-do-i-convert-third-party-types +struct adl_serializer { + static executorch::runtime::Tag from_json(const json &j) { + auto s = j.get(); + if (s == "Tensor") { + return executorch::runtime::Tag::Tensor; + } + if (s == "Int") { + return executorch::runtime::Tag::Int; + } + if (s == "None") { + return executorch::runtime::Tag::None; + } + if (s == "Bool") { + return executorch::runtime::Tag::Bool; + } + if (s == "Double") { + return executorch::runtime::Tag::Double; + } + if (s == "String") { + return executorch::runtime::Tag::String; + } + if (s == "ListInt") { + return executorch::runtime::Tag::ListInt; + } + if (s == "ListBool") { + return executorch::runtime::Tag::ListBool; + } + if (s == "ListDouble") { + return executorch::runtime::Tag::ListDouble; + } + if (s == "ListTensor") { + return executorch::runtime::Tag::ListTensor; + } + throw std::runtime_error(std::format("unknown param kind '{}'", s)); + } + static void to_json(json &j, executorch::runtime::Tag t) { + j = executorch::runtime::tag_to_string(t); + } +}; +} // namespace nlohmann + +namespace rnexecutorch::core::schema { +namespace types = rnexecutorch::core::types; + +using executorch::runtime::Tag; +using nlohmann::json; + +namespace { + +template +struct overloaded : Ts... { + using Ts::operator()...; +}; + +template +T unwrap(const std::string &ctx, executorch::runtime::Result result) { + if (!result.ok()) { + throw std::runtime_error(std::format("{}: {}", ctx, executorch::runtime::to_string(result.error()))); + } + return std::move(result.get()); +} + +} // namespace + +// ======================================================== +// Nlohmann JSON conversions +// ======================================================== + +NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(RangeDim, min, max, step) +NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(DimRef, paramSide, tensorIdx, dimIdx) +NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(EqualityConstraint, dims) +NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(LinearConstraint, dimLhs, dimRhs, coefficients) +NLOHMANN_JSON_SERIALIZE_ENUM(ParamSide, {{ParamSide::input, "input"}, {ParamSide::output, "output"}}) + +// NOLINTNEXTLINE(misc-use-internal-linkage): ADL requires external linkage. +void from_json(const json &j, ConcreteDim &d) { + auto kind = j.at("kind").get(); + if (kind == "constant") { + d = j.at("value").get(); + } else if (kind == "range") { + d = j.at("range").get(); + } else if (kind == "enum") { + d = EnumDim{.choices = j.at("choices").get>()}; + } else { + throw std::runtime_error(std::format("unsupported dim kind '{}'", kind)); + } +} +// NOLINTNEXTLINE(misc-use-internal-linkage): ADL requires external linkage. +void to_json(json &j, const ConcreteDim &d) { + // clang-format off + std::visit(overloaded{ + [&](int32_t c) { j = json::object({{"kind", "constant"}, {"value", c}}); }, + [&](const RangeDim &r) { j = json::object({{"kind", "range"}, {"range", r}}); }, + [&](const EnumDim &e) { j = json::object({{"kind", "enum"}, {"choices", e.choices}}); }, + }, d); + // clang-format on +} + +// NOLINTNEXTLINE(misc-use-internal-linkage): ADL requires external linkage. +void from_json(const json &j, ParamSpec &p) { + p.tag = j.at("kind").get(); + if (p.tag == Tag::Tensor) { + p.dtype = types::parseDType(j.at("dtype").get()); + p.shape = j.at("shape").get>(); + } +} +// NOLINTNEXTLINE(misc-use-internal-linkage): ADL requires external linkage. +void to_json(json &j, const ParamSpec &p) { + if (p.tag == Tag::Tensor) { + // DType is (de)serialized via its string helpers — a JSON macro for it + // would have to live in namespace `types` for ADL to find it. + j = json::object({{"kind", "Tensor"}, {"dtype", types::toString(p.dtype)}, {"shape", p.shape}}); + } else { + j = json::object({{"kind", p.tag}}); + } +} + +// NOLINTNEXTLINE(misc-use-internal-linkage): ADL requires external linkage. +void from_json(const json &j, RuntimeConstraint &c) { + auto kind = j.at("kind").get(); + if (kind == "equality") { + c = j.get(); + } else if (kind == "linear") { + c = j.get(); + } else { + throw std::runtime_error(std::format("unknown constraint kind '{}'", kind)); + } +} +// NOLINTNEXTLINE(misc-use-internal-linkage): ADL requires external linkage. +void to_json(json &j, const RuntimeConstraint &c) { + // clang-format off + std::visit(overloaded{ + [&](const EqualityConstraint &eq) { j = json::object({{"kind", "equality"}, {"dims", eq.dims}}); }, + [&](const LinearConstraint &lin) { + j = json::object({{"kind", "linear"}, + {"dimLhs", lin.dimLhs}, + {"dimRhs", lin.dimRhs}, + {"coefficients", lin.coefficients}}); + }, + }, c); + // clang-format on +} + +NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(MethodSpec, inputs, outputs, runtimeConstraints) + +// ======================================================== +// Serialization +// ======================================================== + +ModelSpec parseModelSpecJson(const std::string &ctx, const std::string &jsonStr) { + try { + return json::parse(jsonStr).get(); + } catch (const std::exception &e) { + throw std::runtime_error(std::format("{}: {}", ctx, e.what())); + } +} + +namespace { + +// NOLINTNEXTLINE(misc-no-recursion): recursion is bounded by JSON nesting depth. +jsi::Value jsonToJs(jsi::Runtime &rt, const json &j) { + switch (j.type()) { + case json::value_t::null: + return jsi::Value::null(); + case json::value_t::boolean: + return jsi::Value(j.get()); + case json::value_t::number_integer: + case json::value_t::number_unsigned: + return jsi::Value(static_cast(j.get())); + case json::value_t::number_float: + return jsi::Value(j.get()); + case json::value_t::string: + return jsi::Value(jsi::String::createFromUtf8(rt, j.get())); + case json::value_t::array: { + auto arr = jsi::Array(rt, j.size()); + size_t i = 0; + for (const auto &el : j) { + arr.setValueAtIndex(rt, i++, jsonToJs(rt, el)); + } + return arr; + } + case json::value_t::object: { + auto obj = jsi::Object(rt); + for (const auto &[k, v] : j.items()) { + obj.setProperty(rt, k.c_str(), jsonToJs(rt, v)); + } + return obj; + } + default: + return jsi::Value::undefined(); + } +} + +} // namespace + +jsi::Value modelSpecToJs(jsi::Runtime &rt, const ModelSpec &spec) { + return jsonToJs(rt, spec); +} + +jsi::Object backendsToJs(jsi::Runtime &rt, + const std::unordered_map> &backends) { + jsi::Object obj(rt); + for (const auto &[methodName, backendList] : backends) { + auto arr = jsi::Array(rt, backendList.size()); + for (size_t i = 0; i < backendList.size(); ++i) { + arr.setValueAtIndex(rt, i, jsi::String::createFromUtf8(rt, backendList[i])); + } + obj.setProperty(rt, methodName.c_str(), arr); + } + return obj; +} + +// ======================================================== +// Metadata reflection +// ======================================================== + +namespace { + +ParamSpec tensorMetaToParamSpec(const executorch::runtime::TensorInfo &tensorMeta) { + const auto sizes = tensorMeta.sizes(); + return ParamSpec{ + .tag = Tag::Tensor, + .dtype = types::fromScalarType(tensorMeta.scalar_type()), + .shape = std::vector(sizes.begin(), sizes.end()), + }; +} + +} // namespace + +MethodSpec methodSpecFromMetadata(const executorch::runtime::MethodMeta &methodMeta) { + MethodSpec spec; + + spec.inputs.reserve(methodMeta.num_inputs()); + for (size_t i = 0; i < methodMeta.num_inputs(); ++i) { + auto ctx = std::format("buildFromMetadata input[{}]", i); + auto tag = unwrap(ctx, methodMeta.input_tag(i)); + if (tag == Tag::Tensor) { + auto tensorMeta = unwrap(ctx, methodMeta.input_tensor_meta(i)); + spec.inputs.emplace_back(tensorMetaToParamSpec(tensorMeta)); + } else { + spec.inputs.emplace_back(ParamSpec{.tag = tag}); + } + } + + spec.outputs.reserve(methodMeta.num_outputs()); + for (size_t i = 0; i < methodMeta.num_outputs(); ++i) { + auto ctx = std::format("buildFromMetadata output[{}]", i); + auto tag = unwrap(ctx, methodMeta.output_tag(i)); + if (tag == Tag::Tensor) { + auto tensorMeta = unwrap(ctx, methodMeta.output_tensor_meta(i)); + spec.outputs.emplace_back(tensorMetaToParamSpec(tensorMeta)); + } else { + spec.outputs.emplace_back(ParamSpec{.tag = tag}); + } + } + + return spec; +} + +std::vector getUsedBackends(const executorch::runtime::MethodMeta &methodMeta) { + std::vector backends; + for (size_t i = 0; i < methodMeta.num_backends(); ++i) { + auto ctx = std::format("getUsedBackends: backend [{}]", i); + const auto *name = unwrap(ctx, methodMeta.get_backend_name(i)); + if (methodMeta.uses_backend(name) && std::ranges::find(backends, name) == backends.end()) { + backends.emplace_back(name); + } + } + return backends; +} + +// ======================================================== +// Validation +// ======================================================== + +namespace { + +void validateConcreteDim(const ConcreteDim &dim, const std::string &ctx) { + // clang-format off + std::visit(overloaded{ + [&](int32_t c) { + if (c <= 0) { + throw std::runtime_error(std::format("{}: constant dim must be positive", ctx)); + } + }, + [&](const RangeDim &r) { + if (r.min < 0) { + throw std::runtime_error(std::format("{}: range min must be non-negative", ctx)); + } + if (r.max < r.min) { + throw std::runtime_error(std::format("{}: range max must be >= min", ctx)); + } + if (r.step <= 0) { + throw std::runtime_error(std::format("{}: range step must be positive", ctx)); + } + }, + [&](const EnumDim &e) { + if (e.choices.empty()) { + throw std::runtime_error(std::format("{}: enum must have at least one choice", ctx)); + } + for (const auto &choice : e.choices) { + if (choice <= 0) { + throw std::runtime_error(std::format("{}: enum choices must be positive", ctx)); + } + } + }, + }, dim); + // clang-format on +} + +void validateSpecDimDomains(const MethodSpec &spec, const std::string &ctx) { + auto validateParams = [&](const std::vector ¶ms, const char *label) { + for (size_t i = 0; i < params.size(); ++i) { + if (params[i].tag != Tag::Tensor) { + continue; + } + for (size_t d = 0; d < params[i].shape.size(); ++d) { + auto dctx = std::format("{} {}[{}] dim[{}]", ctx, label, i, d); + validateConcreteDim(params[i].shape[d], dctx); + } + } + }; + validateParams(spec.inputs, "input"); + validateParams(spec.outputs, "output"); +} + +void validateTensorParam(const ParamSpec ¶m, + const executorch::runtime::TensorInfo &tensorMeta, + const std::string &ctx) { + auto metaDtype = types::fromScalarType(tensorMeta.scalar_type()); + if (param.dtype != metaDtype) { + throw std::runtime_error(std::format("{}: dtype mismatch (spec type '{}' != compiled metadata type '{}')", + ctx, types::toString(param.dtype), types::toString(metaDtype))); + } + + auto metaShape = tensorMeta.sizes(); + if (param.shape.size() != metaShape.size()) { + throw std::runtime_error(std::format("{}: rank mismatch (spec rank {} != compiled metadata rank {})", + ctx, param.shape.size(), metaShape.size())); + } + + for (size_t d = 0; d < param.shape.size(); ++d) { + auto bound = static_cast(metaShape[d]); + // clang-format off + std::visit(overloaded{ + [&](int32_t c) { + if (c != bound) { + throw std::runtime_error(std::format("{}: shape[{}] mismatch (spec constant {} != compiled bound {})", + ctx, d, c, bound)); + } + }, + [&](const RangeDim &r) { + if (r.max > bound) { + throw std::runtime_error(std::format("{}: shape[{}] range max {} exceeds compiled bound {}", + ctx, d, r.max, bound)); + } + }, + [&](const EnumDim &e) { + for (const auto choice : e.choices) { + if (choice > bound) { + throw std::runtime_error(std::format("{}: shape[{}] enum choice {} exceeds compiled bound {}", + ctx, d, choice, bound)); + } + } + }, + }, + param.shape[d]); + // clang-format on + } +} + +void validateDimRef(const DimRef &ref, + const std::vector &inputRanks, + const std::vector &outputRanks, + const std::string &ctx) { + bool isInput = (ref.paramSide == ParamSide::input); + const auto &ranks = isInput ? inputRanks : outputRanks; + if (std::cmp_greater_equal(ref.tensorIdx, ranks.size())) { + throw std::runtime_error(std::format("{}: tensorIdx {} out of range", ctx, ref.tensorIdx)); + } + if (std::cmp_greater_equal(ref.dimIdx, ranks[static_cast(ref.tensorIdx)])) { + throw std::runtime_error(std::format("{}: dimIdx {} out of range", ctx, ref.dimIdx)); + } +} + +std::vector gatherTensorRanks(const std::vector ¶ms) { + std::vector ranks; + for (const auto &p : params) { + if (p.tag == Tag::Tensor) { + ranks.push_back(p.shape.size()); + } + } + return ranks; +} + +void validateConstraintSpecs(const MethodSpec &spec, const std::string &ctx) { + auto inputRanks = gatherTensorRanks(spec.inputs); + auto outputRanks = gatherTensorRanks(spec.outputs); + + for (size_t i = 0; i < spec.runtimeConstraints.size(); ++i) { + auto cctx = std::format("{} constraint[{}]", ctx, i); + const auto &constraint = spec.runtimeConstraints[i]; + + // clang-format off + std::visit(overloaded{ + [&](const EqualityConstraint &eq) { + if (eq.dims.size() < 2) { + throw std::runtime_error(std::format("{}: equality needs at least two dims", cctx)); + } + for (const auto &dim : eq.dims) { + validateDimRef(dim, inputRanks, outputRanks, cctx); + } + }, + [&](const LinearConstraint &lin) { + validateDimRef(lin.dimLhs, inputRanks, outputRanks, cctx); + validateDimRef(lin.dimRhs, inputRanks, outputRanks, cctx); + }, + }, constraint); + // clang-format on + } +} + +void validateParamsAgainstMeta(const std::vector ¶ms, + bool isInput, + const executorch::runtime::MethodMeta &meta, + const std::string &ctx) { + for (size_t i = 0; i < params.size(); ++i) { + auto pctx = std::format("{} {}[{}]", ctx, isInput ? "input" : "output", i); + auto tagResult = isInput ? unwrap(pctx, meta.input_tag(i)) + : unwrap(pctx, meta.output_tag(i)); + + if (params[i].tag != tagResult) { + throw std::runtime_error(std::format("{}: tag mismatch (spec tag {} != compiled metadata tag {})", + pctx, executorch::runtime::tag_to_string(params[i].tag), + executorch::runtime::tag_to_string(tagResult))); + } + + if (tagResult == Tag::Tensor) { + auto tensorMeta = isInput ? unwrap(pctx, meta.input_tensor_meta(i)) + : unwrap(pctx, meta.output_tensor_meta(i)); + validateTensorParam(params[i], tensorMeta, pctx); + } + } +} + +} // namespace + +void validateSpec(const MethodSpec &spec, + const executorch::runtime::MethodMeta &meta, + const std::string &ctx) { + + if (spec.inputs.size() != meta.num_inputs()) { + throw std::runtime_error(std::format("{}: input count mismatch (spec has {}, model metadata has {})", + ctx, spec.inputs.size(), meta.num_inputs())); + } + if (spec.outputs.size() != meta.num_outputs()) { + throw std::runtime_error(std::format("{}: output count mismatch (spec has {}, model metadata has {})", + ctx, spec.outputs.size(), meta.num_outputs())); + } + + validateSpecDimDomains(spec, ctx); + validateParamsAgainstMeta(spec.inputs, /*isInput=*/true, meta, ctx); + validateParamsAgainstMeta(spec.outputs, /*isInput=*/false, meta, ctx); + validateConstraintSpecs(spec, ctx); +} + +namespace { + +int32_t getInputDimValue(const DimRef &ref, + const std::vector> &inputShapes) { + return inputShapes[static_cast(ref.tensorIdx)][static_cast(ref.dimIdx)]; +} + +} // namespace + +void validateRuntimeConstraints(jsi::Runtime &rt, + const std::vector &constraints, + const std::vector> &inputShapes, + const std::string &ctx) { + for (size_t i = 0; i < constraints.size(); ++i) { + auto cctx = std::format("{} constraint[{}]", ctx, i); + + // clang-format off + std::visit(overloaded{ + [&](const EqualityConstraint &eq) { + std::vector inputVals; + for (const auto &d : eq.dims) { + if (d.paramSide == ParamSide::input) { + inputVals.push_back(getInputDimValue(d, inputShapes)); + } + } + if (inputVals.size() < 2) { + return; + } + for (size_t j = 1; j < inputVals.size(); ++j) { + if (inputVals[j] != inputVals[0]) { + throw jsi::JSError(rt, std::format("{}: equality constraint violated (dimension value {} != {})", + cctx, inputVals[0], inputVals[j])); + } + } + }, + [&](const LinearConstraint &lin) { + if (lin.dimLhs.paramSide == ParamSide::output || + lin.dimRhs.paramSide == ParamSide::output) { + return; + } + int32_t lhs = getInputDimValue(lin.dimLhs, inputShapes); + int32_t rhs = getInputDimValue(lin.dimRhs, inputShapes); + if (lhs != lin.coefficients[0] * rhs + lin.coefficients[1]) { + throw jsi::JSError(rt, std::format("{}: linear constraint violated (LHS {} != {} * RHS {} + {})", + cctx, lhs, lin.coefficients[0], rhs, lin.coefficients[1])); + } + }, + }, constraints[i]); + // clang-format on + } +} + +} // namespace rnexecutorch::core::schema diff --git a/packages/react-native-executorch/cpp/core/schema.h b/packages/react-native-executorch/cpp/core/schema.h new file mode 100644 index 0000000000..484108eb65 --- /dev/null +++ b/packages/react-native-executorch/cpp/core/schema.h @@ -0,0 +1,289 @@ +#pragma once + +#include "core/dtype.h" +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +/** + * Model spec types, JSON parsing/serialization, and validation. + * + * This module is the C++ counterpart of `src/core/schema.ts`. Both files + * define the exact same structural data model (`ModelSpec`, `MethodSpec`, + * `ParamSpec`, `ConcreteDim`, `RuntimeConstraint`, etc.) and validation + * semantics, but operate at different phases of the lifecycle: + * + * - **schema.ts (TypeScript)** — Validates *allowed* (symbolic) specs written by + * pipeline authors against *exported* (concrete) specs. Handles symbol binding, + * dimension domain matching, and 1-to-1 constraint matching. Runs in JavaScript/TypeScript. + * + * - **schema.h / schema.cpp (Native C++)** — Parses exported JSON specs, validates + * them at model load time against ExecuTorch `MethodMeta`, and performs dynamic + * shape and constraint validation prior to model execution. Runs in native C++. + + * Model specifications originate from one of two sources at load time: + * 1. **Optional JSON Companion Method**: When a `.pte` model binary exports the + * companion method `rnexecutorch::core::model::kGetModelSchemaMethod`, + * `ModelHostObject` calls it to retrieve a JSON string. This string is + * parsed by `parseModelSpecJson` into a `ModelSpec` containing rich dynamic + * dimension domains (`RangeDim`, `EnumDim`) and `RuntimeConstraint`s. + * 2. **Fallback MethodMeta**: If companion is absent, `methodSpecFromMetadata` + * derives a basic `ModelSpec` directly from static ExecuTorch `MethodMeta` + * (where all dimensions are static `int32_t` constants and no constraints + * are present). + * + * The types defined in `schema.h` mirror their TypeScript counterparts in + * `src/core/schema.ts`: + + * Validation in C++ occurs in two distinct phases: + * + * 1. **Load-Time Spec Validation (`validateSpec`)**: + * Executed inside `ModelHostObject` constructor when loading a model. It checks + * the parsed exported spec (`MethodSpec`) against ExecuTorch's `MethodMeta` metadata. + * Verifies that parameter counts, primitive tags, tensor `DType`s, and static shape + * dimensions match ExecuTorch's compiled program requirements. + * + * 2. **Dynamic Runtime Validation (`validateRuntimeConstraints`)**: + * Executed natively in C++ inside `ModelHostObject::execute` immediately prior to + * running model inference. It inspects the actual concrete dimensions of user-provided + * input tensors and asserts that: + * - Dynamic dimensions fall within their declared domains (`RangeDim` + * min/max/step or `EnumDim` choices). + * - `EqualityConstraint`: Referenced input tensor dimensions evaluate to + * identical integer values. + * - `LinearConstraint`: Referenced input dimensions satisfy `dimLhs == a * + * dimRhs + b`. + * + * *Note on Output Dimension References*: Pre-execution validation only checks + * constraints over input dimensions (`ParamSide::input`). References to output + * dimensions (`ParamSide::output`) are skipped prior to execution (for `EqualityConstraint`, + * output dimension references are filtered out; for `LinearConstraint`, constraints referencing + * an output dimension are skipped). After execution, concrete output shapes produced by + * ExecuTorch are verified when copying data into user-provided output buffers. + * + * @see src/core/schema.ts for the TypeScript validation layer. + * @see rnexecutorch::core::model::kGetModelSchemaMethod for the companion method constant. + */ +namespace rnexecutorch::core::schema { +namespace jsi = facebook::jsi; + +// ======================================================== +// Parameter specs +// ======================================================== + +/** + * Inclusive integer domain of a single dynamic dimension — values from `min` + * to `max` in increments of `step`. + */ +struct RangeDim { + int32_t min = 0; + int32_t max = 0; + int32_t step = 1; +}; + +/** + * A single dimension matching one of the listed `choices`. + */ +struct EnumDim { + std::vector choices; +}; + +/** + * A dimension of an exported (concrete) model spec. + */ +using ConcreteDim = std::variant; + +/** + * A single input or output parameter of a method. `tag` is the discriminator: + * `executorch::runtime::Tag::Tensor` for tensor params (in which case + * `dtype`/`shape` describe the tensor), or any primitive tag for non-tensor + * params (in which case `dtype`/`shape` are unused). + */ +struct ParamSpec { + executorch::runtime::Tag tag = executorch::runtime::Tag::None; + types::DType dtype{}; ///< Valid iff `tag == Tag::Tensor`. + std::vector shape; ///< Valid iff `tag == Tag::Tensor`. +}; + +// ======================================================== +// Runtime constraints +// ======================================================== + +/// Whether the referenced tensor dimension belongs to an input or output. +enum class ParamSide { input, + output }; + +/** + * Reference to a single tensor dimension of a method's input or output. + * `tensorIdx` counts only tensor parameters (skipping primitives), consistent + * with ExecuTorch's `inputTensorMeta` / `outputTensorMeta` ordering. + */ +struct DimRef { + ParamSide paramSide = ParamSide::input; + int32_t tensorIdx = 0; + int32_t dimIdx = 0; +}; + +/** + * Runtime constraint declaring that all referenced dimensions must be equal + * to each other in any given execution of the method. + */ +struct EqualityConstraint { + std::vector dims; +}; + +/** + * Runtime constraint declaring that two dimensions must satisfy + * `dimLhs = coefficients[0] * dimRhs + coefficients[1]` (integer coefficients) + * in any given execution of the method. + */ +struct LinearConstraint { + DimRef dimLhs{}; + DimRef dimRhs{}; + std::array coefficients{}; +}; + +/** + * A requirement on the runtime values of a method's tensor dimensions: the + * concrete tensors passed to and produced by the method must satisfy it in + * any given execution. + */ +using RuntimeConstraint = std::variant; + +// ======================================================== +// Model specs +// ======================================================== + +/** + * Spec of a single model method: the ordered input and output parameter + * specs and the runtime constraints the method declares over its tensor + * dimensions. + */ +struct MethodSpec { + std::vector inputs; + std::vector outputs; + std::vector runtimeConstraints; +}; + +/** + * Spec of a whole model, mapping method names to their MethodSpec. + */ +using ModelSpec = std::unordered_map; + +// ======================================================== +// Serialization +// ======================================================== + +/** + * Parses a JSON-encoded ModelSpec into a C++ `ModelSpec` structure using + * nlohmann::json. Throws `std::runtime_error` with `ctx` context on invalid + * JSON syntax or unrecognized kinds and tags. Semantic validation against model + * metadata is performed separately by `validateSpec`. + * + * @param ctx Context description used for error messages. + * @param json The JSON string to parse. + * @return The parsed model spec. + */ +ModelSpec parseModelSpecJson(const std::string &ctx, const std::string &json); + +/** + * Serializes a ModelSpec to a fresh JS object matching ModelSpec + * from src/core/schema.ts. + * + * @param rt The JSI runtime instance. + * @param spec The model spec to serialize. + * @return The spec as a JS value. + */ +jsi::Value modelSpecToJs(jsi::Runtime &rt, const ModelSpec &spec); + +/** + * Serializes a backends map to a JSI object mapping method names to arrays of + * backend name strings. + * + * @param rt The JSI runtime instance. + * @param backends Map of method name -> backend name list. + * @return The backends as a JSI object. + */ +jsi::Object backendsToJs(jsi::Runtime &rt, + const std::unordered_map> &backends); + +// ======================================================== +// Metadata reflection +// ======================================================== + +/** + * Builds a MethodSpec directly from ExecuTorch MethodMeta, without requiring + * a JSON companion. All shapes become ConcreteDim constants since MethodMeta + * only exposes the static shape from export. + * + * @param methodMeta The method metadata from the .pte program. + * @return A MethodSpec derived from the metadata. + */ +MethodSpec methodSpecFromMetadata(const executorch::runtime::MethodMeta &methodMeta); + +/** + * Collects the names of backends declared as used by a method. + * + * @param methodMeta The method metadata from the .pte program. + * @return Vector of backend names for which uses_backend returns true. + */ +std::vector getUsedBackends(const executorch::runtime::MethodMeta &methodMeta); + +// ======================================================== +// Validation +// ======================================================== + +/** + * Validates a MethodSpec against ExecuTorch MethodMeta at load time. Performs + * the following checks: + * 1. Dimension domain validity: positive constant values, ordered range bounds + * (`min > 0`, `max >= min`, `step > 0`), and non-empty positive enum + * choices. + * 2. Parameter metadata matching: input/output counts, parameter primitive + * tags, tensor `DType`s, tensor ranks, exact values of static constant + * dimensions, and dynamic dimension upper bounds (`RangeDim` max and + * `EnumDim` choices <= compiled `MethodMeta` allocation bound). + * 3. Constraint structure: verifies `EqualityConstraint` has at least 2 + * dimensions and all `DimRef` tensor and dimension indices are within valid + * input/output rank bounds. + * + * @param spec The method spec to validate. + * @param meta The method metadata from the .pte program. + * @param ctx Context string for error messages. + * @throws std::runtime_error on any mismatch or invalid spec. + */ +void validateSpec(const MethodSpec &spec, + const executorch::runtime::MethodMeta &meta, + const std::string &ctx); + +/** + * Validates runtime constraints against actual concrete input tensor shapes before execution. + * + * References to output dimensions (`ParamSide::output`) are skipped prior to + * execution (for `EqualityConstraint`, output dimension references are filtered + * out; for `LinearConstraint`, constraints referencing an output dimension are + * skipped). After execution, concrete output shapes produced by ExecuTorch are + * verified when copying data into user-provided output buffers. + * + * @param rt The JSI runtime instance (for error reporting). + * @param constraints The list of runtime constraints to validate. + * @param inputShapes Per-tensor input shapes (indexed by `DimRef.tensorIdx`). + * @param ctx Context string for error messages (e.g. method name). + * @throws jsi::JSError on constraint violation. + */ +void validateRuntimeConstraints(jsi::Runtime &rt, + const std::vector &constraints, + const std::vector> &inputShapes, + const std::string &ctx); + +} // namespace rnexecutorch::core::schema diff --git a/packages/react-native-executorch/cpp/core/tensor_helpers.cpp b/packages/react-native-executorch/cpp/core/tensor_helpers.cpp index 54c9e23064..de77106ae2 100644 --- a/packages/react-native-executorch/cpp/core/tensor_helpers.cpp +++ b/packages/react-native-executorch/cpp/core/tensor_helpers.cpp @@ -1,8 +1,14 @@ #include "tensor_helpers.h" +#include +#include #include +#include +#include #include +#include +#include "core/schema.h" #include "dtype.h" namespace rnexecutorch::core::tensor { @@ -10,69 +16,82 @@ namespace types = rnexecutorch::core::types; namespace conversions = rnexecutorch::core::conversions; std::shared_lock -tryLockShared(jsi::Runtime &rt, const std::string &name, const std::shared_ptr &tensor) { +tryLockShared(jsi::Runtime &rt, const std::string &ctx, const std::shared_ptr &tensor) { std::shared_lock lock(tensor->mutex_, std::try_to_lock); if (!lock.owns_lock()) { - throw jsi::JSError(rt, std::format("{} tensor is currently in use", name)); + throw jsi::JSError(rt, std::format("{} tensor is currently in use", ctx)); } if (!tensor->data_) { - throw jsi::JSError(rt, std::format("{} tensor has been disposed", name)); + throw jsi::JSError(rt, std::format("{} tensor has been disposed", ctx)); } return lock; } std::unique_lock -tryLockUnique(jsi::Runtime &rt, const std::string &name, const std::shared_ptr &tensor) { +tryLockUnique(jsi::Runtime &rt, const std::string &ctx, const std::shared_ptr &tensor) { std::unique_lock lock(tensor->mutex_, std::try_to_lock); if (!lock.owns_lock()) { - throw jsi::JSError(rt, std::format("{} tensor is currently in use", name)); + throw jsi::JSError(rt, std::format("{} tensor is currently in use", ctx)); } if (!tensor->data_) { - throw jsi::JSError(rt, std::format("{} tensor has been disposed", name)); + throw jsi::JSError(rt, std::format("{} tensor has been disposed", ctx)); } return lock; } void checkNotSameTensor(jsi::Runtime &rt, - const std::string &name1, const std::shared_ptr &t1, - const std::string &name2, const std::shared_ptr &t2) { + const std::string &ctx1, const std::shared_ptr &t1, + const std::string &ctx2, const std::shared_ptr &t2) { if (t1 == t2) { - throw jsi::JSError(rt, std::format("{} and {} cannot be the same tensor", name1, name2)); + throw jsi::JSError(rt, std::format("{} and {} cannot be the same tensor", ctx1, ctx2)); } } namespace { + +template +struct overloaded : Ts... { + using Ts::operator()...; +}; + std::string shapeToString(const SymbolicShape &shape) { std::string s; for (const auto &dim : shape) { - if (std::holds_alternative(dim)) { - s += std::to_string(std::get(dim)); - } else if (std::holds_alternative(dim)) { - s += std::get(dim); - } else { - const auto &rangeDim = std::get(dim); - s += std::format("[{}..{}{}]", - rangeDim.min, - rangeDim.max, - rangeDim.step ? std::format(" step {}", *rangeDim.step) : ""); - } - s += ", "; + // clang-format off + std::visit(overloaded{ + [&](const std::string &str) { s += str; }, + [&](int32_t val) { s += std::to_string(val); }, + [&](const schema::RangeDim &range) { + s += std::format("[{}..{}:{}]", range.min, range.max, range.step); + }, + [&](const schema::EnumDim &enumeration) { + s += "{"; + for (const auto choice : enumeration.choices) { + s += std::to_string(choice) + ","; + } + if (!enumeration.choices.empty()) { + s.pop_back(); + } + s += "}"; + }, + }, dim); + // clang-format on + s += ","; } if (!shape.empty()) { s.pop_back(); - s.pop_back(); } return "[" + s + "]"; } } // namespace std::shared_ptr -fromJs(jsi::Runtime &rt, const std::string &name, const jsi::Value &value, +fromJs(jsi::Runtime &rt, const std::string &ctx, const jsi::Value &value, std::optional expectedDtype, const std::optional &expectedShape) { - auto obj = conversions::asType(rt, name, value); + auto obj = conversions::asType(rt, ctx, value); if (!obj.isHostObject(rt)) { - throw jsi::JSError(rt, name + " must be a Tensor"); + throw jsi::JSError(rt, ctx + " must be a Tensor"); } auto tensor = obj.getHostObject(rt); @@ -80,7 +99,7 @@ fromJs(jsi::Runtime &rt, const std::string &name, const jsi::Value &value, const auto &shape = tensor->shape_; if (expectedDtype && dtype != *expectedDtype) { - throw jsi::JSError(rt, std::format("{} must be of type {}", name, types::toString(*expectedDtype))); + throw jsi::JSError(rt, std::format("{} must be of type {} (got {})", ctx, types::toString(*expectedDtype), types::toString(dtype))); } if (!expectedShape) { @@ -89,42 +108,51 @@ fromJs(jsi::Runtime &rt, const std::string &name, const jsi::Value &value, if (shape.size() != expectedShape->size()) { throw jsi::JSError(rt, std::format("{} must have shape {} (expected {} dimensions, got {})", - name, shapeToString(*expectedShape), expectedShape->size(), shape.size())); + ctx, shapeToString(*expectedShape), expectedShape->size(), shape.size())); } - std::unordered_map symbolToConcrete; + std::unordered_map symbolBinding; for (size_t i = 0; i < expectedShape->size(); ++i) { const auto &dim = expectedShape->at(i); - if (std::holds_alternative(dim)) { - const auto expected = std::get(dim); - if (shape[i] != expected) { - throw jsi::JSError(rt, std::format("{} must have shape {} (dim {} mismatch: expected {}, got {})", - name, shapeToString(*expectedShape), i, expected, shape[i])); - } - } else if (std::holds_alternative(dim)) { - const auto &symbol = std::get(dim); - if (symbolToConcrete.contains(symbol) && shape[i] != symbolToConcrete[symbol]) { - throw jsi::JSError(rt, std::format("{} must have shape {} (dim {} mismatch: expected {}, got {})", - name, shapeToString(*expectedShape), i, symbolToConcrete[symbol], shape[i])); - } - symbolToConcrete[symbol] = shape[i]; - } else { - const auto &rangeDim = std::get(dim); - if (shape[i] < rangeDim.min) { - throw jsi::JSError(rt, std::format("{} must have shape {} (dim {} out of range: {} < min {})", - name, shapeToString(*expectedShape), i, shape[i], rangeDim.min)); - } - if (shape[i] > rangeDim.max) { - throw jsi::JSError(rt, std::format("{} must have shape {} (dim {} out of range: {} > max {})", - name, shapeToString(*expectedShape), i, shape[i], rangeDim.max)); - } - if (rangeDim.step && (shape[i] - rangeDim.min) % *rangeDim.step != 0) { - throw jsi::JSError(rt, std::format("{} must have shape {} (dim {} must be min({}) + k*step({}), got {})", - name, shapeToString(*expectedShape), i, rangeDim.min, *rangeDim.step, shape[i])); - } - } + // clang-format off + std::visit(overloaded{ + [&](const std::string &symbol) { + if (symbolBinding.contains(symbol) && symbolBinding[symbol] != shape[i]) { + throw jsi::JSError(rt, std::format("{} must have shape {} (symbol {} mismatch: expected {}, got {})", + ctx, shapeToString(*expectedShape), symbol, symbolBinding[symbol], shape[i])); + } + symbolBinding[symbol] = shape[i]; + }, + [&](int32_t val) { + if (shape[i] != val) { + throw jsi::JSError(rt, std::format("{} must have shape {} (dim {} mismatch: expected {}, got {})", + ctx, shapeToString(*expectedShape), i, val, shape[i])); + } + }, + [&](const schema::RangeDim &range) { + if (shape[i] < range.min) { + throw jsi::JSError(rt, std::format("{} must have shape {} (dim {} out of range: {} < min {})", + ctx, shapeToString(*expectedShape), i, shape[i], range.min)); + } + if (shape[i] > range.max) { + throw jsi::JSError(rt, std::format("{} must have shape {} (dim {} out of range: {} > max {})", + ctx, shapeToString(*expectedShape), i, shape[i], range.max)); + } + if ((shape[i] - range.min) % range.step != 0) { + throw jsi::JSError(rt, std::format("{} must have shape {} (dim {} must be min({}) + k*step({}), got {})", + ctx, shapeToString(*expectedShape), i, range.min, range.step, shape[i])); + } + }, + [&](const schema::EnumDim &enumeration) { + if (std::ranges::find(enumeration.choices, shape[i]) == enumeration.choices.end()) { + throw jsi::JSError(rt, std::format("{} must have shape {} (dim {} not allowed: got {})", + ctx, shapeToString(*expectedShape), i, shape[i])); + } + }, + }, dim); + // clang-format on } return tensor; diff --git a/packages/react-native-executorch/cpp/core/tensor_helpers.h b/packages/react-native-executorch/cpp/core/tensor_helpers.h index 9004e6f262..2c616d7276 100644 --- a/packages/react-native-executorch/cpp/core/tensor_helpers.h +++ b/packages/react-native-executorch/cpp/core/tensor_helpers.h @@ -14,6 +14,7 @@ #include "conversions.h" #include "dtype.h" +#include "schema.h" #include "tensor.h" #include @@ -22,14 +23,11 @@ namespace rnexecutorch::core::tensor { namespace jsi = facebook::jsi; using rnexecutorch::core::types::DType; +using schema::EnumDim; +using schema::RangeDim; -struct RangeDim { - int32_t min = 0; - int32_t max = 0; - std::optional step; -}; - -using SymbolicShape = std::vector>; +using SymbolicDim = std::variant; +using SymbolicShape = std::vector; /** * Tries to acquire a shared (read) lock on the underlying tensor resource. @@ -42,7 +40,7 @@ using SymbolicShape = std::vector>; * @return A shared lock protecting the tensor data. */ [[nodiscard]] std::shared_lock -tryLockShared(jsi::Runtime &rt, const std::string &name, const std::shared_ptr &tensor); +tryLockShared(jsi::Runtime &rt, const std::string &ctx, const std::shared_ptr &tensor); /** * Tries to acquire a unique (write) lock on the underlying tensor resource. @@ -50,12 +48,12 @@ tryLockShared(jsi::Runtime &rt, const std::string &name, const std::shared_ptr -tryLockUnique(jsi::Runtime &rt, const std::string &name, const std::shared_ptr &tensor); +tryLockUnique(jsi::Runtime &rt, const std::string &ctx, const std::shared_ptr &tensor); /** * Validates that two JSI Tensor parameters do not point to the exact same @@ -63,14 +61,14 @@ tryLockUnique(jsi::Runtime &rt, const std::string &name, const std::shared_ptr &t1, - const std::string &name2, const std::shared_ptr &t2); + const std::string &ctx1, const std::shared_ptr &t1, + const std::string &ctx2, const std::shared_ptr &t2); /** * Extracts, type-checks, and shape-validates a TensorHostObject from a JSI @@ -78,16 +76,30 @@ void checkNotSameTensor(jsi::Runtime &rt, * checking or shape validation fails. * * @param rt The JSI runtime instance. - * @param name Parameter name for contextual error messages. + * @param ctx Parameter name for contextual error messages. * @param value The JSI value to extract the Tensor from. * @param expectedDtype Optional expected DType constraint. * @param expectedShape Optional expected shape (SymbolicShape) constraint. * @return Shared pointer to the validated TensorHostObject. */ std::shared_ptr -fromJs(jsi::Runtime &rt, const std::string &name, const jsi::Value &value, +fromJs(jsi::Runtime &rt, const std::string &ctx, const jsi::Value &value, std::optional expectedDtype, const std::optional &expectedShape); +/** + * @overload + * + * Convenience wrapper that accepts an initializer list of symbolic shape + * elements. Allows passing shape constraints like `{"H", "W", 1}` directly + * without typing SymbolicShape explicitly. + */ +inline std::shared_ptr +fromJs(jsi::Runtime &rt, const std::string &ctx, const jsi::Value &value, + std::optional expectedDtype, + std::initializer_list expectedShape) { + return fromJs(rt, ctx, value, expectedDtype, std::optional(expectedShape)); +} + /** * @overload * @@ -100,23 +112,28 @@ template requires std::ranges::input_range && std::convertible_to, int32_t> inline std::shared_ptr -fromJs(jsi::Runtime &rt, const std::string &name, const jsi::Value &value, +fromJs(jsi::Runtime &rt, const std::string &ctx, const jsi::Value &value, std::optional expectedDtype, const Range &expectedShape) { SymbolicShape convertedShape(expectedShape.begin(), expectedShape.end()); - return fromJs(rt, name, value, expectedDtype, std::move(convertedShape)); + return fromJs(rt, ctx, value, expectedDtype, std::move(convertedShape)); } /** * @overload * - * Convenience wrapper that accepts an initializer list of symbolic shape - * elements. Allows passing shape constraints like `{"H", "W", 1}` directly - * without typing SymbolicShape explicitly. + * Convenience wrapper that accepts a vector of ConcreteDim values as the + * expected shape. Useful when building shapes programmatically from + * ConcreteDim without manually wrapping each element in SymbolicDim. */ inline std::shared_ptr -fromJs(jsi::Runtime &rt, const std::string &name, const jsi::Value &value, - std::optional expectedDtype, - std::initializer_list> expectedShape) { - return fromJs(rt, name, value, expectedDtype, std::optional(expectedShape)); +fromJs(jsi::Runtime &rt, const std::string &ctx, const jsi::Value &value, + std::optional expectedDtype, const std::vector &expectedShape) { + SymbolicShape convertedShape; + convertedShape.reserve(expectedShape.size()); + for (const auto &dim : expectedShape) { + convertedShape.push_back(std::visit([](const auto &d) -> SymbolicDim { return d; }, dim)); + } + return fromJs(rt, ctx, value, expectedDtype, std::move(convertedShape)); } + } // namespace rnexecutorch::core::tensor diff --git a/packages/react-native-executorch/src/core/model.ts b/packages/react-native-executorch/src/core/model.ts index e539afecdd..5fd59bb2d4 100644 --- a/packages/react-native-executorch/src/core/model.ts +++ b/packages/react-native-executorch/src/core/model.ts @@ -1,5 +1,6 @@ import { rnexecutorchJsi } from '../native/bridge'; -import type { DType, Tensor } from './tensor'; +import type { Tensor } from './tensor'; +import type { ModelSpec, ConcreteDim } from './schema'; declare const modelBrand: unique symbol; @@ -13,67 +14,7 @@ export type ModelInput = Tensor | number | boolean | null; * A value returned from a model's `execute` method. * @category Types */ -export type ModelOutput = Tensor | number | boolean | null; - -/** - * Metadata describing a single tensor slot (input or output) of a model method. - * @category Types - */ -export type TensorMeta = { - /** The name associated with this tensor slot (may be empty). */ - name: string; - /** The number of dimensions. */ - ndim: number; - /** The total byte size of the tensor buffer. */ - nbytes: number; - /** The element data type. */ - dtype: DType; - /** The concrete size of each dimension (e.g. `[1, 3, 224, 224]`). */ - shape: number[]; -}; - -/** - * The ExecuTorch value-tag that classifies the runtime type of a model input or - * output slot. - * @category Types - */ -export type ExecuTorchTag = - | 'None' - | 'Tensor' - | 'Int' - | 'Double' - | 'Bool' - | 'String' - | 'ListBool' - | 'ListDouble' - | 'ListInt' - | 'ListTensor'; - -/** - * Metadata describing a single exported method of an ExecuTorch model. - * @category Types - */ -export type ModelMethodMeta = { - /** The exported method name (e.g. `'forward'`). */ - name: string; - /** The total number of input arguments the method accepts. */ - numInputs: number; - /** The total number of output values the method returns. */ - numOutputs: number; - /** Runtime value-tags for each input slot, in order. */ - inputTags: ExecuTorchTag[]; - /** Runtime value-tags for each output slot, in order. */ - outputTags: ExecuTorchTag[]; - /** - * A map from backend name to a boolean indicating whether this method - * delegates to that backend. - */ - usesBackend: Record; - /** Detailed tensor metadata for every input tensor slot, in order. */ - inputTensorMeta: TensorMeta[]; - /** Detailed tensor metadata for every output tensor slot, in order. */ - outputTensorMeta: TensorMeta[]; -}; +export type ModelOutput = Tensor | number | boolean | string | null; /** * A compiled, ready-to-run ExecuTorch model loaded into native memory. @@ -89,20 +30,10 @@ export type ModelMethodMeta = { export interface Model { /** The local filesystem path of the `.pte` model file. */ readonly path: string; - - /** - * Returns the list of exported method names available on this model (e.g. - * `['forward']`). - */ - getMethodNames(): string[]; - - /** - * Returns detailed metadata for the specified exported method, including - * input/output tags, tensor shapes, dtype, and backend delegation info. - * @param methodName The name of the exported method to inspect. - * @returns The {@link ModelMethodMeta} for the requested method. - */ - getMethodMeta(methodName: string): ModelMethodMeta; + /** The exported schema of this model. */ + readonly schema: ModelSpec; + /** ExecuTorch backends used by a given method. */ + readonly backends: Record; /** * Executes a named model method synchronously. diff --git a/packages/react-native-executorch/src/core/modelSchema.ts b/packages/react-native-executorch/src/core/modelSchema.ts deleted file mode 100644 index 0fcdc579ef..0000000000 --- a/packages/react-native-executorch/src/core/modelSchema.ts +++ /dev/null @@ -1,223 +0,0 @@ -import { type DType } from './tensor'; -import { type Model, type ExecuTorchTag, type ModelMethodMeta, type TensorMeta } from './model'; - -/** - * A single dimension in a symbolic tensor shape. - * - * - A **number** matches only that exact dimension size (e.g. `1`, `3`). - * - A **string** is a symbolic variable that can match any integer and is - * resolved consistently within the same tensor shape (e.g. `'N'`, `'H'`, - * `'W'`). - * @category Types - */ -export type SymbolicShape = readonly (number | string)[]; - -/** - * A constraint on the dtype and/or shape of a tensor input or output slot. - * - * Both fields are optional; omitting `dtype` skips dtype checking, and omitting - * `shapes` (or providing an empty array) skips shape checking. - * @category Types - */ -export type TensorConstraint = { - /** If set, the tensor's dtype must match exactly. */ - readonly dtype?: DType; - /** - * One or more acceptable symbolic shapes. The tensor passes validation if - * its concrete shape matches **at least one** of the listed shapes. - */ - readonly shapes?: readonly SymbolicShape[]; -}; - -/** - * Convenience constructor for a {@link TensorConstraint}. - * - * ```ts - * // Accepts a float32 tensor with shape [1, 3, H, W] or [3, H, W] - * SymbolicTensor('float32', [1, 3, 'H', 'W'], [3, 'H', 'W']) - * - * // Accepts any tensor of shape [1, N] - * SymbolicTensor(undefined, [1, 'N']) - * ``` - * @category Typescript API - * @param dtype Optional dtype requirement. Pass `undefined` to skip dtype - * checking. - * @param shapes Zero or more acceptable symbolic shapes. An empty list skips - * shape checking. - * @returns A {@link TensorConstraint} object. - */ -export function SymbolicTensor( - dtype?: DType, - ...shapes: readonly SymbolicShape[] -): TensorConstraint { - return { dtype, shapes }; -} - -const primitiveTagMap = { - number: ['Int', 'Double'] as ExecuTorchTag[], - boolean: ['Bool'] as ExecuTorchTag[], - null: ['None'] as ExecuTorchTag[], -} as const; - -/** - * A constraint describing an expected input or output slot of a model method. - * - * - `'number'` — the slot must carry an integer or double primitive value. - * - `'boolean'` — the slot must carry a boolean primitive value. - * - `'null'` — the slot must carry a `None` value. - * - {@link TensorConstraint} — the slot must carry a tensor, optionally with a - * constrained dtype and/or shape. - * @category Types - */ -export type ValueConstraint = keyof typeof primitiveTagMap | TensorConstraint; - -/** - * Checks whether a concrete tensor shape matches at least one of the provided - * symbolic shapes. - * - * Symbolic dimensions (strings) act as named wildcards: within a single shape - * candidate they must resolve to the same integer value every time they appear, - * but they are not enforced consistently across different tensor slots. - * @category Typescript API - * @param actual The concrete shape array to test (e.g. `[1, 3, 224, 224]`). - * @param expected One or more symbolic shapes to match against. - * @returns `true` if `actual` matches at least one of the `expected` shapes. - */ -export function matchShape(actual: number[], ...expected: readonly SymbolicShape[]): boolean { - return expected.some((shape) => { - if (actual.length !== shape.length) return false; - const symbolMap = new Map(); - return shape.every((dim, i) => { - const act = actual[i]!; - if (typeof dim === 'number') return act === dim; - if (symbolMap.has(dim)) return symbolMap.get(dim) === act; - symbolMap.set(dim, act); - return true; - }); - }); -} - -// TODO: Implement cross-tensor symbol validation (e.g. enforcing that 'N' or -// 'H' has the same value across different input/output tensors). Note that a -// proper implementation will require backtracking/solving when multiple tensors -// have multiple alternative shapes. For now, we just check that each tensor -// individually matches at least one of its expected shapes, without enforcing -// consistency of symbolic dimensions across tensors, only within each tensor. -function validateTags( - side: 'input' | 'output', - expected: readonly ValueConstraint[], - actualTags: ExecuTorchTag[], - tensorMetas: TensorMeta[] -) { - const numTensors = expected.filter((t) => typeof t === 'object').length; - if (tensorMetas.length !== numTensors) - throw new Error( - `${side} tensor count mismatch: expected ${numTensors}, got ${tensorMetas.length}` - ); - - let tIdx = 0; - expected.forEach((exp, i) => { - const act = actualTags[i]!; - if (typeof exp === 'string') { - if (!primitiveTagMap[exp].includes(act)) { - throw new Error(`${side}[${i}]: expected primitive '${exp}', got '${act}'`); - } - } else { - if (act !== 'Tensor') { - throw new Error(`${side}[${i}]: expected Tensor, got primitive '${act}'`); - } - const tMeta = tensorMetas[tIdx++]!; - if (exp.dtype && tMeta.dtype !== exp.dtype) { - throw new Error( - `${side}[${i}]: dtype mismatch: expected '${exp.dtype}', got '${tMeta.dtype}'` - ); - } - if (exp.shapes?.length && !matchShape(tMeta.shape, ...exp.shapes)) { - const expectedShapesStr = exp.shapes.map((s) => `[${s.join(',')}]`).join('|'); - throw new Error( - `${side}[${i}]: shape mismatch: expected shape matching ${expectedShapesStr}, got [${tMeta.shape.join(',')}]` - ); - } - } - }); -} - -/** - * Validates that a compiled model's method signature matches the declared - * input and output constraints, throwing a descriptive error on mismatch. - * - * The function checks: - * - That the method exists on the model. - * - That the number of input and output slots matches the expected counts. - * - That the value-tag of each slot (tensor, int, bool, etc.) is compatible - * with its declared {@link ValueConstraint}. - * - For tensor slots: that the dtype and shape (if specified) satisfy the - * - * {@link TensorConstraint}. - * - * On success it returns the method's {@link ModelMethodMeta}, which can be used - * to read concrete input/output tensor shapes for pre-allocating scratch - * tensors. - * @remarks - * A symbolic (string) input dimension here only relaxes *validation*. For a - * dimension that must genuinely vary at runtime (e.g. sequence length), the - * `.pte` must additionally be exported with a companion method - * `get_dynamic_dims_` (e.g. `get_dynamic_dims_forward`) that - * returns, per tensor input, an `int32` `[rank, 3]` tensor of `[min, max, step]` - * bounds — ExecuTorch metadata only serializes the static upper bound, so the - * runtime reads the active range from this companion. Without it, the method - * only accepts the exact shape it was exported with. The reported upper bound is - * still read from `meta.inputTensorMeta`. - * @category Typescript API - * @param model The compiled model to validate. - * @param methodName The exported method name to validate (e.g. `'forward'`). - * @param expectedInputs Ordered list of {@link ValueConstraint}s for each input - * slot. - * @param expectedOutputs Ordered list of {@link ValueConstraint}s for each - * output slot. - * @returns The {@link ModelMethodMeta} for the validated method. - * @throws {Error} A human-readable description of which constraint failed. - */ -export function validateModelSchema( - model: Model, - methodName: string, - expectedInputs: readonly ValueConstraint[], - expectedOutputs: readonly ValueConstraint[] -): ModelMethodMeta { - if (!model.getMethodNames().includes(methodName)) - throw new Error(`signature validation: '${methodName}' method not found`); - - const meta = model.getMethodMeta(methodName); - - const formatError = (errorMsg: string) => { - return ( - `signature validation failed for '${methodName}': ${errorMsg}\n` + - ` Expected: ${JSON.stringify(expectedInputs)} -> ${JSON.stringify(expectedOutputs)}\n` + - ` Actual: [${meta.inputTags.join(', ')}] -> [${meta.outputTags.join(', ')}] (metas: ${JSON.stringify(meta.inputTensorMeta)} -> ${JSON.stringify(meta.outputTensorMeta)})` - ); - }; - - if (meta.inputTags.length !== expectedInputs.length) { - throw new Error( - formatError( - `input count mismatch: expected ${expectedInputs.length}, got ${meta.inputTags.length}` - ) - ); - } - if (meta.outputTags.length !== expectedOutputs.length) { - throw new Error( - formatError( - `output count mismatch: expected ${expectedOutputs.length}, got ${meta.outputTags.length}` - ) - ); - } - - try { - validateTags('input', expectedInputs, meta.inputTags, meta.inputTensorMeta); - validateTags('output', expectedOutputs, meta.outputTags, meta.outputTensorMeta); - } catch (e: any) { - throw new Error(formatError(e.message)); - } - - return meta; -} diff --git a/packages/react-native-executorch/src/core/schema.ts b/packages/react-native-executorch/src/core/schema.ts new file mode 100644 index 0000000000..4ef0325041 --- /dev/null +++ b/packages/react-native-executorch/src/core/schema.ts @@ -0,0 +1,824 @@ +/** + * Model specs and spec validation. + * + * A model spec is a structural contract describing a model's methods: the + * parameter specs of every input and output (primitive tags, tensor data + * types, and per-dimension domains) and the runtime constraints the method + * declares over its tensor dimensions. A spec is either: + * - **allowed** (`SymbolicDim`) — written by a pipeline to state which models + * it can work with. Dimensions may be named symbols: `static` symbols bind + * to constant dimensions, `dynamic` symbols to ranges or enums, and reusing + * a symbol requires every occurrence to bind to the same domain. Several + * allowed specs can be passed as variants; matching any one of them is + * enough. + * - **exported** (`ConcreteDim`) — derived from an exported model's metadata, + * stating what the model actually provides. + * + * **Dimension domains vs runtime values.** The central distinction of this + * module is between a dimension's *domain* and its *runtime value*. The + * domain is the set of values the dimension may take: `constant` is a + * singleton (the value is fully known statically), while `range` and `enum` + * are proper sets that only narrow the possibilities. The runtime value is + * the actual size of the dimension for a concrete tensor in one given + * execution — a single element drawn from the domain. + * + * Accordingly, two different validations must not be confused: + * - **Spec validation** (this module, {@link validateSpec}) — a static, + * load-time check that an exported spec satisfies an allowed spec. It only + * ever compares domains: symbols bind to domains (repeated symbols to the + * same one), and constraints are matched as declarations. Domain equality + * says nothing about runtime values — two dimensions with the same domain + * may still take different values in an execution. + * - **Runtime validation** (the native runtime, not this module) — at + * execution time the runtime checks that every concrete tensor shape lies + * within the declared domains and enforces the declared runtime + * constraints. + * + * Runtime constraints are statements about runtime values, not domains: an + * equality constraint requires its dimensions' values to coincide in any + * given execution, and a linear constraint requires them to satisfy + * `lhs = a * rhs + b`. Since spec validation only sees domains, it cannot + * decide whether such a relation holds — the exported spec must simply + * declare exactly the allowed spec's constraints (1-to-1, no missing, no + * extras). The only exception is degenerate: a constant domain has exactly + * one possible value, so an equality constraint between constants is fully + * decided statically — equal constants are equal at runtime. Linear + * constraints, in contrast, are never evaluated against domains here, even + * between constants. + * + * **Exported spec source.** An exported model's `ModelSpec` + * is populated at load time from one of two sources: + * + * 1. **ExecuTorch `MethodMeta`** (default) — when the `.pte` only carries + * static metadata, every dimension domain is `constant`. This is + * sufficient for models whose input/output shapes are fully fixed at + * export time. + * + * 2. **Companion `get_model_schema` method** — for models whose tensors + * have dynamic or enumerated dimensions (e.g. variable-length sequences) + * or that declare runtime constraints, the `.pte` exports a method named + * `get_model_schema` that returns a JSON-encoded `ModelSpec` + * string. The native loader calls this method after loading the model and + * merges the result into `model.schema`, overlaying precise `range`, + * `enum`, and `RuntimeConstraint` entries onto the base `MethodMeta`. + * Only methods that actually need overrides need to appear in the JSON; + * methods absent from the companion spec are kept as-is from + * `MethodMeta`. + * + * **Tip for model export:** + * Embed the companion schema method during ExecuTorch compilation in Python + * by passing `constant_methods={"get_model_schema": schema_json}` when + * lowering with `to_edge_transform_and_lower(...)` where `schema_json` is + * the JSON string encoding the model's `ModelSpec`. + * @packageDocumentation + */ +import type { DType } from './tensor'; + +// ======================================================== +// Parameter specs +// ======================================================== + +/** + * Inclusive integer domain of a single dynamic dimension — values from `min` + * to `max` in increments of `step`. + * @category Types + */ +export type Range = { readonly min: number; readonly max: number; readonly step: number }; + +/** + * A single dimension with a fully known domain: + * - `constant` — exactly `value`. + * - `range` — any value of a {@link Range}. + * - `enum` — one of the listed `choices`. + * @category Types + */ +export type ConcreteDim = + | { readonly kind: 'constant'; readonly value: number } + | { readonly kind: 'range'; readonly range: Range } + | { readonly kind: 'enum'; readonly choices: readonly number[] }; + +/** + * A single dimension of an allowed model spec. On top of {@link ConcreteDim} + * domains, a named symbol binds to the exported spec's dimension at + * validation: `static` symbols bind to constants, `dynamic` symbols to ranges + * or enums. Reusing a symbol requires every occurrence to bind to the same + * domain — it does NOT imply any runtime relation between the dimensions. + * @category Types + */ +export type SymbolicDim = + | ConcreteDim + | { readonly kind: 'static'; readonly symbol: string } + | { readonly kind: 'dynamic'; readonly symbol: string }; + +/** + * Spec of a tensor parameter: the expected element `dtype` and one + * dimension spec per axis. + * @category Types + */ +export type TensorSpec = { + readonly kind: 'Tensor'; + readonly dtype: DType; + readonly shape: readonly Dim[]; +}; + +/** + * The ExecuTorch value-tag that classifies the runtime type of a model input or + * output slot. + * @category Types + */ +export type ExecuTorchTag = + | 'None' + | 'Tensor' + | 'Int' + | 'Double' + | 'Bool' + | 'String' + | 'ListBool' + | 'ListDouble' + | 'ListInt' + | 'ListTensor'; + +/** + * Spec of a single input or output parameter of a method — either a + * {@link TensorSpec} or a primitive ExecuTorch value tag (`Int`, `Bool`, ...). + * @category Types + */ +export type ParamSpec = + | TensorSpec + | { readonly kind: Exclude }; + +// ======================================================== +// Runtime constraints +// ======================================================== + +/** + * Reference to a single tensor dimension of a method's input or output. + * `tensorIdx` counts only tensor parameters (skipping primitives), consistent + * with ExecuTorch's `inputTensorMeta` / `outputTensorMeta` ordering. + * @category Types + */ +export type DimRef = { + readonly paramSide: 'input' | 'output'; + readonly tensorIdx: number; + readonly dimIdx: number; +}; + +/** + * Runtime constraint declaring that all referenced dimensions must be equal + * to each other in any given execution of the method. + * @category Types + */ +export type EqualityConstraint = { + readonly kind: 'equality'; + readonly dims: readonly DimRef[]; +}; + +/** + * Runtime constraint declaring that two dimensions must satisfy + * `dimLhs = coefficients[0] * dimRhs + coefficients[1]` (integer + * coefficients) in any given execution of the method. + * @category Types + */ +export type LinearConstraint = { + readonly kind: 'linear'; + readonly dimLhs: DimRef; + readonly dimRhs: DimRef; + readonly coefficients: [number, number]; +}; + +/** + * A requirement on the runtime values of a method's tensor dimensions: the + * concrete tensors passed to and produced by the method must satisfy it in + * any given execution. Matched as a declaration during spec validation. + * @category Types + */ +export type RuntimeConstraint = LinearConstraint | EqualityConstraint; + +// ======================================================== +// Model specs +// ======================================================== + +/** + * Spec of a single model method: the ordered input and output parameter specs + * and the runtime constraints the method declares over its tensor dimensions. + * @category Types + */ +export type MethodSpec = { + inputs: readonly ParamSpec[]; + outputs: readonly ParamSpec[]; + runtimeConstraints: readonly RuntimeConstraint[]; +}; + +/** + * Spec of a whole model, mapping method names to their {@link MethodSpec}. + * A `SymbolicDim` spec describes allowed models; a `ConcreteDim` spec + * describes an exported model. + * @category Types + */ +export type ModelSpec = Record>; + +// ======================================================== +// Helper functions +// ======================================================== + +/** + * Shape notation accepted by {@link SymbolicTensor}: numbers become + * {@link ConstantDim}, strings become {@link StaticDim}, and + * {@link SymbolicDim} values are used as-is. + * @category Types + */ +export type SymbolicShape = readonly (number | string | SymbolicDim)[]; + +/** + * Creates a static symbolic dimension. Static symbols bind to constant + * dimensions of the exported spec; repeated uses must bind to the same value. + * @category Typescript API + * @param symbol The symbol name. + * @returns The symbolic dimension. + */ +export const StaticDim = (symbol: string): SymbolicDim => { + return { kind: 'static', symbol }; +}; + +/** + * Creates a dynamic symbolic dimension. Dynamic symbols bind to range or enum + * dimensions of the exported spec; repeated uses must bind to the same domain. + * @category Typescript API + * @param symbol The symbol name. + * @returns The symbolic dimension. + */ +export const DynamicDim = (symbol: string): SymbolicDim => { + return { kind: 'dynamic', symbol }; +}; + +/** + * Creates a constant dimension matching exactly `value`. + * @category Typescript API + * @param value The required dimension size. + * @returns The concrete dimension. + * @throws {Error} If `value` is not a positive integer. + */ +export const ConstantDim = (value: number): ConcreteDim => { + if (value <= 0 || !Number.isInteger(value)) { + throw new Error(`Invalid value (${value}): must be a positive integer.`); + } + return { kind: 'constant', value }; +}; + +/** + * Creates an enumerated dimension matching one of `choices`. + * @category Typescript API + * @param choices The allowed dimension sizes. + * @returns The concrete dimension. + * @throws {Error} If any choice is not a positive integer. + */ +export const EnumDim = (choices: readonly number[]): ConcreteDim => { + if (choices.some((dim) => dim <= 0 || !Number.isInteger(dim))) { + throw new Error(`Invalid enum choice: must be a positive integer`); + } + return { kind: 'enum', choices }; +}; + +/** + * Creates a range dimension matching values from `min` to `max` in increments + * of `step`. + * @category Typescript API + * @param min The smallest allowed dimension size. + * @param max The largest allowed dimension size. + * @param step The increment between allowed sizes. Defaults to 1. + * @returns The concrete dimension. + * @throws {Error} If the range bounds or step are not valid positive integers. + */ +export const RangeDim = (min: number, max: number, step?: number): ConcreteDim => { + if (min < 0 || !Number.isInteger(min)) { + throw new Error(`Invalid range min (${min}): must be a non-negative integer.`); + } + if (max < min) { + throw new Error(`Invalid range [${min}, ${max}]: max cannot be less than min.`); + } + if (!Number.isInteger(max)) { + throw new Error(`Invalid range max (${max}): must be a non-negative integer.`); + } + if (step !== undefined && (step <= 0 || !Number.isInteger(step))) { + throw new Error(`Invalid range step (${step}): must be a positive integer.`); + } + return { kind: 'range', range: { min, max, step: step ?? 1 } }; +}; + +/** + * Creates a {@link TensorSpec} from a dtype and a {@link SymbolicShape}: + * numbers become {@link ConstantDim}, strings become {@link StaticDim}. + * @category Typescript API + * @param dtype The expected element data type. + * @param shape The per-dimension specs. + * @returns The tensor spec. + */ +export const SymbolicTensor = (dtype: DType, shape: SymbolicShape) => { + const typedShape = shape.map((dim) => { + if (typeof dim === 'string') return StaticDim(dim); + if (typeof dim === 'number') return ConstantDim(dim); + return dim; + }); + return { kind: 'Tensor', dtype, shape: typedShape } as TensorSpec; +}; + +export const f32 = (...shape: SymbolicShape) => SymbolicTensor('float32', shape); +export const i64 = (...shape: SymbolicShape) => SymbolicTensor('int64', shape); +export const i32 = (...shape: SymbolicShape) => SymbolicTensor('int32', shape); +export const ui8 = (...shape: SymbolicShape) => SymbolicTensor('uint8', shape); + +/** Helper namespace for declaring runtime constraints. */ +export const constr = { + eq: (...dims: DimRef[]): EqualityConstraint => { + return { kind: 'equality', dims }; + }, + linear: (dimLhs: DimRef, dimRhs: DimRef, a: number, b: number = 0): LinearConstraint => { + return { kind: 'linear', dimLhs, dimRhs, coefficients: [a, b] }; + }, +}; + +/** + * Constructs a method specification mapping a method name to its parameter + * specs and runtime constraints. + * @param name The execution method name (e.g., `'forward'`). + * @param inputs Ordered list of parameter specifications for method inputs. + * @param outputs Ordered list of parameter specifications for method outputs. + * @param constraints Optional array of {@link RuntimeConstraint} declarations + * (equality or linear relations across dimensions). + * @returns A record mapping `name` to its corresponding {@link MethodSpec}. + */ +export function method( + name: string, + inputs: ParamSpec[], + outputs: ParamSpec[], + constraints?: RuntimeConstraint[] +): Record> { + return { [name]: { inputs, outputs, runtimeConstraints: constraints ?? [] } }; +} + +// ======================================================== +// Parameter spec validation +// ======================================================== + +type SymbolBindings = Map; + +function rangesEqual(r1: Range, r2: Range): boolean { + return r1.min === r2.min && r1.max === r2.max && r1.step === r2.step; +} + +function choicesEqual(c1: readonly number[], c2: readonly number[]): boolean { + const s1 = new Set(c1); + const s2 = new Set(c2); + return s1.size === s2.size && [...s1].every((elem) => s2.has(elem)); +} + +function matchDim( + sDim: SymbolicDim, + cDim: ConcreteDim, + bindings: SymbolBindings, + ctx: string +): void { + if (sDim.kind === 'constant' && cDim.kind === 'constant') { + if (sDim.value !== cDim.value) { + throw new Error(`${ctx}: Constant dimension mismatch.`); + } + return; + } + + if (sDim.kind === 'range' && cDim.kind === 'range') { + if (!rangesEqual(sDim.range, cDim.range)) { + throw new Error(`${ctx}: Range dimension mismatch.`); + } + return; + } + + if (sDim.kind === 'enum' && cDim.kind === 'enum') { + if (!choicesEqual(sDim.choices, cDim.choices)) { + throw new Error(`${ctx}: Enum dimension mismatch.`); + } + return; + } + + if (sDim.kind === 'static' && cDim.kind === 'constant') { + const bind = bindings.get(sDim.symbol); + if (bind) { + if (bind.kind !== 'constant' || bind.value !== cDim.value) { + throw new Error(`${ctx}: Symbol '${sDim.symbol}' has inconsistent bindings.`); + } + return; + } + bindings.set(sDim.symbol, cDim); + return; + } + + if (sDim.kind === 'dynamic' && (cDim.kind === 'range' || cDim.kind === 'enum')) { + const bind = bindings.get(sDim.symbol); + if (bind) { + const consistentRange = + bind.kind === 'range' && cDim.kind === 'range' && rangesEqual(bind.range, cDim.range); + const consistentEnum = + bind.kind === 'enum' && cDim.kind === 'enum' && choicesEqual(bind.choices, cDim.choices); + if (!consistentRange && !consistentEnum) { + throw new Error(`${ctx}: Symbol '${sDim.symbol}' has inconsistent bindings.`); + } + return; + } + bindings.set(sDim.symbol, cDim); + return; + } + + throw new Error(`${ctx}: Cannot match symbolic '${sDim.kind}' with concrete '${cDim.kind}'.`); +} + +function matchMethodSpecs( + allowedMethodSpec: MethodSpec, + exportedMethodSpec: MethodSpec, + bindings: SymbolBindings, + ctx: string +): void { + if (allowedMethodSpec.inputs.length !== exportedMethodSpec.inputs.length) { + throw new Error(`${ctx}: Input count mismatch.`); + } + if (allowedMethodSpec.outputs.length !== exportedMethodSpec.outputs.length) { + throw new Error(`${ctx}: Output count mismatch.`); + } + + const allowedParamSpecs = [...allowedMethodSpec.inputs, ...allowedMethodSpec.outputs]; + const exportedParamSpecs = [...exportedMethodSpec.inputs, ...exportedMethodSpec.outputs]; + + for (let p = 0; p < allowedParamSpecs.length; ++p) { + const isInput = p < allowedMethodSpec.inputs.length; + const paramSpecIdx = isInput ? p : p - allowedMethodSpec.inputs.length; + const paramSpecCtx = `${ctx} ${isInput ? 'input' : 'output'} #${paramSpecIdx}`; + + const allowedParamSpec = allowedParamSpecs[p]!; + const exportedParamSpec = exportedParamSpecs[p]!; + + if (allowedParamSpec.kind !== exportedParamSpec.kind) { + throw new Error(`${paramSpecCtx}: Param spec kind mismatch.`); + } + + if (allowedParamSpec.kind !== 'Tensor') continue; + + const allowedTensorSpec = allowedParamSpec as TensorSpec; + const exportedTensorSpec = exportedParamSpec as TensorSpec; + + if (allowedTensorSpec.dtype !== exportedTensorSpec.dtype) { + throw new Error(`${paramSpecCtx}: DType mismatch.`); + } + + if (allowedTensorSpec.shape.length !== exportedTensorSpec.shape.length) { + throw new Error(`${paramSpecCtx}: Rank mismatch.`); + } + + for (let d = 0; d < allowedTensorSpec.shape.length; ++d) { + const dimCtx = `${paramSpecCtx} Tensor dim #${d}`; + matchDim(allowedTensorSpec.shape[d]!, exportedTensorSpec.shape[d]!, bindings, dimCtx); + } + } +} + +function matchModelSpecsSymbols( + allowedModelSpec: ModelSpec, + exportedModelSpec: ModelSpec, + bindings: SymbolBindings +): void { + for (const [methodName, allowedMethodSpec] of Object.entries(allowedModelSpec)) { + const exportedMethodSpec = exportedModelSpec[methodName]; + if (!exportedMethodSpec) { + throw new Error(`Method '${methodName}' not found in exported model spec.`); + } + matchMethodSpecs(allowedMethodSpec, exportedMethodSpec, bindings, `Method '${methodName}'`); + } +} + +// ======================================================== +// Runtime constraints validation +// ======================================================== + +function refsEqual(r1: DimRef, r2: DimRef): boolean { + return r1.paramSide === r2.paramSide && r1.tensorIdx === r2.tensorIdx && r1.dimIdx === r2.dimIdx; +} + +function constraintsEqual(c1: RuntimeConstraint, c2: RuntimeConstraint): boolean { + if (c1.kind === 'linear' && c2.kind === 'linear') { + return ( + refsEqual(c1.dimLhs, c2.dimLhs) && + refsEqual(c1.dimRhs, c2.dimRhs) && + c1.coefficients[0] === c2.coefficients[0] && + c1.coefficients[1] === c2.coefficients[1] + ); + } + + if (c1.kind === 'equality' && c2.kind === 'equality') { + if (c1.dims.length !== c2.dims.length) { + return false; + } + + const unclaimed = [...c2.dims]; + for (const ref of c1.dims) { + const idx = unclaimed.findIndex((r) => refsEqual(r, ref)); + if (idx === -1) return false; + unclaimed.splice(idx, 1); + } + return true; + } + + return false; +} + +function resolveDim(methodSpec: MethodSpec, ref: DimRef): D { + let tensorSpecs: TensorSpec[]; + switch (ref.paramSide) { + case 'input': + tensorSpecs = methodSpec.inputs.filter((v): v is TensorSpec => v.kind === 'Tensor'); + break; + case 'output': + tensorSpecs = methodSpec.outputs.filter((v): v is TensorSpec => v.kind === 'Tensor'); + break; + } + + const tensorSpec = tensorSpecs[ref.tensorIdx]; + if (!tensorSpec) { + throw new Error(`Invalid DimRef (${JSON.stringify(ref)}): tensor index out of range.`); + } + + const dim = tensorSpec.shape[ref.dimIdx]; + if (!dim) { + throw new Error(`Invalid DimRef (${JSON.stringify(ref)}): dimension index out of range.`); + } + + return dim; +} + +function matchRuntimeConstraints( + allowedModelSpec: ModelSpec, + exportedModelSpec: ModelSpec +): void { + for (const [methodName, allowedMethodSpec] of Object.entries(allowedModelSpec)) { + const exportedMethodSpec = exportedModelSpec[methodName]; + if (!exportedMethodSpec) { + throw new Error(`Method '${methodName}' not found in exported model spec.`); + } + + const unclaimed = [...exportedMethodSpec.runtimeConstraints]; + + for (const [idx, constraint] of allowedMethodSpec.runtimeConstraints.entries()) { + const find = unclaimed.findIndex((c) => constraintsEqual(c, constraint)); + if (find === -1) { + throw new Error(`Constraint ${idx}: Not declared by the exported model spec.`); + } + unclaimed.splice(find, 1); + } + + if (unclaimed.length > 0) { + throw new Error(`'${methodName}': Exported spec declares unexpected runtime constraints`); + } + } +} + +// ======================================================== +// Spec validation +// ======================================================== + +function validateSymbolKindConsistency(modelSpec: ModelSpec): void { + const symbolKinds = new Map(); + + for (const methodSpec of Object.values(modelSpec)) { + for (const paramSpec of [...methodSpec.inputs, ...methodSpec.outputs]) { + if (paramSpec.kind !== 'Tensor') continue; + + for (const dim of paramSpec.shape) { + if (!('symbol' in dim)) continue; + + const existing = symbolKinds.get(dim.symbol); + if (existing && existing !== dim.kind) { + throw new Error(`Invalid spec: '${dim.symbol}' is used as both 'static' and 'dynamic'.`); + } + symbolKinds.set(dim.symbol, dim.kind); + } + } + } +} + +function validateConstraintCorrectness(modelSpec: ModelSpec): void { + for (const [methodName, methodSpec] of Object.entries(modelSpec)) { + for (const [idx, constraint] of methodSpec.runtimeConstraints.entries()) { + const ctx = `Method '${methodName}' constraint ${idx}`; + + if (constraint.kind === 'linear') { + const [A, B] = constraint.coefficients; + if (!Number.isInteger(A) || !Number.isInteger(B)) { + throw new Error(`${ctx}: Coefficients must be integers.`); + } + resolveDim(methodSpec, constraint.dimLhs); + resolveDim(methodSpec, constraint.dimRhs); + } + + if (constraint.kind === 'equality') { + if (constraint.dims.length < 2) { + throw new Error(`${ctx}: Equality requires at least two dimensions.`); + } + constraint.dims.forEach((ref) => resolveDim(methodSpec, ref)); + } + } + } +} + +function validateDimDomains(modelSpec: ModelSpec): void { + for (const [methodName, methodSpec] of Object.entries(modelSpec)) { + const allParams = [...methodSpec.inputs, ...methodSpec.outputs]; + for (const [p, param] of allParams.entries()) { + if (param.kind !== 'Tensor') continue; + + const isInput = p < methodSpec.inputs.length; + const label = isInput ? 'input' : 'output'; + const paramIdx = isInput ? p : p - methodSpec.inputs.length; + + for (const [d, dim] of param.shape.entries()) { + const ctx = `Method '${methodName}' ${label} #${paramIdx} dim #${d}`; + + if (dim.kind === 'constant') { + if (dim.value <= 0 || !Number.isInteger(dim.value)) { + throw new Error(`${ctx}: constant dim must be a positive integer.`); + } + } + if (dim.kind === 'range') { + if (dim.range.min < 0 || !Number.isInteger(dim.range.min)) { + throw new Error(`${ctx}: range min must be a non-negative integer.`); + } + if (dim.range.max < dim.range.min || !Number.isInteger(dim.range.max)) { + throw new Error(`${ctx}: range max must be >= min.`); + } + if (dim.range.step <= 0 || !Number.isInteger(dim.range.step)) { + throw new Error(`${ctx}: range step must be a positive integer.`); + } + } + if (dim.kind === 'enum') { + if (dim.choices.length === 0) { + throw new Error(`${ctx}: enum must have at least one choice.`); + } + if (dim.choices.some((c) => c <= 0 || !Number.isInteger(c))) { + throw new Error(`${ctx}: enum choices must be positive integers.`); + } + } + } + } + } +} + +/** + * Result of validating an exported model spec against allowed variants. + * @typeParam K The variant key type. + */ +export type SpecMatch = { + /** Key of the matched variant. */ + readonly variant: K; + + /** + * Returns the concrete value for a symbol. + * @param name The symbol name. + * @param kind Expected dimension kind — determines the return type. Omit to + * get the raw {@link ConcreteDim} when the kind is not known upfront. + * @throws {Error} If the symbol is not found or has a different kind. + */ + dim(name: string, kind: 'constant'): number; + dim(name: string, kind: 'range'): Range; + dim(name: string, kind: 'enum'): readonly number[]; + dim(name: string, kind: 'dynamic'): Exclude; + dim(name: string): ConcreteDim; + + /** + * Batch accessors for retrieving multiple symbol dimensions at once as typed tuples. + */ + dims: { + /** + * Retrieves concrete dimension objects for multiple symbols. + * @param names Symbol names to retrieve. + * @returns A tuple of {@link ConcreteDim} objects corresponding to `names`. + * @throws {Error} If any symbol is not found. + */ + any(...names: S): { [I in keyof S]: ConcreteDim }; + + /** + * Retrieves choice arrays for multiple enumerated dynamic symbols. + * @param names Symbol names expected to be enum dimensions. + * @returns A tuple of choice arrays corresponding to `names`. + * @throws {Error} If any symbol is not found or is not an enum dimension. + */ + enum(...names: S): { [I in keyof S]: readonly number[] }; + + /** + * Retrieves range objects for multiple dynamic symbols. + * @param names Symbol names expected to be range dimensions. + * @returns A tuple of {@link Range} objects corresponding to `names`. + * @throws {Error} If any symbol is not found or is not a range dimension. + */ + range(...names: S): { [I in keyof S]: Range }; + + /** + * Retrieves dynamic dimension objects (range or enum) for multiple symbols. + * @param names Symbol names expected to be dynamic dimensions (range or enum). + * @returns A tuple of dynamic {@link ConcreteDim} objects corresponding to `names`. + * @throws {Error} If any symbol is not found or is a constant dimension. + */ + dynamic( + ...names: S + ): { [I in keyof S]: Exclude }; + + /** + * Retrieves constant numeric values for multiple static symbols. + * @param names Symbol names expected to be constant dimensions. + * @returns A tuple of numbers corresponding to `names`. + * @throws {Error} If any symbol is not found or is not a constant dimension. + */ + constant(...names: S): { [I in keyof S]: number }; + }; +}; + +/** + * Validates that an exported (concrete) model spec satisfies at least one of + * the allowed (symbolic) model specs — variants are tried in order and the + * first match wins. For a variant to match: + * - Every method exists in the exported spec and its signature matches, + * binding each symbol to a constant value (static) or a range/enum + * (dynamic). Repeated symbols must bind consistently across the whole spec. + * - The exported spec declares exactly the same runtime constraints per + * method (1-to-1, no missing, no extras). Constraints are matched as + * declarations only; whether they hold at runtime is the model's guarantee. + * + * Authoring bugs in an allowed spec (conflicting symbol kinds, invalid + * constraint coefficients or references) throw immediately, before matching. + * @param exportedModelSpec The exported model spec to validate against. + * @param allowedModelSpecs The allowed model spec variants keyed by name. + * @returns A {@link SpecMatch} with the matched variant key and dim + * accessors. + * @throws {Error} A human-readable description of why every variant failed. + */ +export function validateSpec>>( + exportedModelSpec: ModelSpec, + allowedModelSpecs: T +): SpecMatch { + const entries = Object.entries(allowedModelSpecs); + + for (const spec of Object.values(allowedModelSpecs)) { + validateSymbolKindConsistency(spec); + validateConstraintCorrectness(spec); + validateDimDomains(spec); + } + validateDimDomains(exportedModelSpec); + validateConstraintCorrectness(exportedModelSpec); + + const errors: string[] = []; + + for (const [key, allowedModelSpec] of entries) { + try { + const bindings: SymbolBindings = new Map(); + matchModelSpecsSymbols(allowedModelSpec, exportedModelSpec, bindings); + matchRuntimeConstraints(allowedModelSpec, exportedModelSpec); + + const dimFn = (name: string, kind?: string): any => { + const dim = bindings.get(name); + if (!dim) { + throw new Error(`Symbol '${name}' not found in bindings.`); + } + if (kind) { + if (kind === 'dynamic') { + if (dim.kind === 'constant') { + throw new Error(`Symbol '${name}' is 'constant', expected 'dynamic'.`); + } + return dim; + } + if (dim.kind !== kind) { + throw new Error(`Symbol '${name}' is '${dim.kind}', expected '${kind}'.`); + } + } + if (dim.kind === 'constant') return dim.value; + if (dim.kind === 'range') return dim.range; + if (dim.kind === 'enum') return dim.choices; + return dim; + }; + + const createAccessor = (kind?: string) => { + return (...names: string[]): any => names.map((name) => dimFn(name, kind)); + }; + + return { + variant: key, + dim: dimFn, + dims: { + any: createAccessor(), + enum: createAccessor('enum'), + range: createAccessor('range'), + dynamic: createAccessor('dynamic'), + constant: createAccessor('constant'), + }, + }; + } catch (e: any) { + errors.push(`Variant '${key}': ${e.message}`); + continue; + } + } + + throw new Error(`Spec doesn't match any of the provided variants:\n - ${errors.join('\n - ')}`); +} diff --git a/packages/react-native-executorch/src/extensions/cv/tasks/classification.ts b/packages/react-native-executorch/src/extensions/cv/tasks/classification.ts index 3a8b2d244c..1c0dd34129 100644 --- a/packages/react-native-executorch/src/extensions/cv/tasks/classification.ts +++ b/packages/react-native-executorch/src/extensions/cv/tasks/classification.ts @@ -2,7 +2,7 @@ import type { WorkletRuntime } from 'react-native-worklets'; import { tensor } from '../../../core/tensor'; import { loadModel } from '../../../core/model'; -import { validateModelSchema, SymbolicTensor } from '../../../core/modelSchema'; +import { validateSpec, method, f32 } from '../../../core/schema'; import { wrapAsync } from '../../../core/runtime'; import { softmax } from '../../math'; @@ -58,6 +58,7 @@ export async function createClassifier( * Releases all allocated native resources. */ dispose: () => void; + /** * Performs asynchronous image classification on the given input image. * @param input The input image buffer. @@ -68,6 +69,7 @@ export async function createClassifier( * confidence. */ classify: (input: ImageBuffer, options?: { topk?: number }) => Promise[]>; + /** * Synchronous version of {@link classify} to be executed directly on the * caller or worklet thread. @@ -77,25 +79,31 @@ export async function createClassifier( const { modelPath, classifierOpts } = config; const model = await wrapAsync(loadModel, runtime)(modelPath); - const meta = validateModelSchema( - model, - 'forward', - [SymbolicTensor('float32', [1, 3, 'H', 'W'], [3, 'H', 'W'])], - [SymbolicTensor('float32', [1, 'N'], ['N'])] - ); - const inpShape = meta.inputTensorMeta[0]!.shape; - const outShape = meta.outputTensorMeta[0]!.shape; - - const numLabels = outShape[outShape.length - 1]!; - if (classifierOpts.labels.length !== numLabels) { + const { variant, dims } = validateSpec(model.schema, { + batched: method( + 'forward', // prettier-ignore + [f32(1, 3, 'H', 'W')], + [f32(1, 'N')] + ), + unbatched: method( + 'forward', // prettier-ignore + [f32(3, 'H', 'W')], + [f32('N')] + ), + }); + + const [N, H, W] = dims.constant('N', 'H', 'W'); + const inpShape = { batched: [1, 3, H, W], unbatched: [3, H, W] }[variant]; + const outShape = { batched: [1, N], unbatched: [N] }[variant]; + + if (classifierOpts.labels.length !== N) { throw new Error( - `Classifier labels length (${classifierOpts.labels.length}) must match model output dimension (${numLabels}).` + `Classifier labels length (${classifierOpts.labels.length}) must match model output dimension (${N}).` ); } - // prettier-ignore const tensors = [ - tensor('float32', outShape), + tensor('float32', outShape), // prettier-ignore tensor('float32', outShape), ] as const; @@ -119,9 +127,8 @@ export async function createClassifier( const tInput = preprocessor.process(input); model.execute('forward', [tInput], [tLogits]); - // prettier-ignore const probas = tLogits - .through(softmax, tProbas) + .through(softmax, tProbas) // prettier-ignore .getData(new Float32Array(tProbas.numel)); return Array.from(probas) diff --git a/packages/react-native-executorch/src/extensions/cv/tasks/imageEmbedding.ts b/packages/react-native-executorch/src/extensions/cv/tasks/imageEmbedding.ts index afd4c63383..22192f5eda 100644 --- a/packages/react-native-executorch/src/extensions/cv/tasks/imageEmbedding.ts +++ b/packages/react-native-executorch/src/extensions/cv/tasks/imageEmbedding.ts @@ -2,7 +2,7 @@ import type { WorkletRuntime } from 'react-native-worklets'; import { tensor } from '../../../core/tensor'; import { loadModel } from '../../../core/model'; -import { validateModelSchema, SymbolicTensor } from '../../../core/modelSchema'; +import { validateSpec, method, f32 } from '../../../core/schema'; import { wrapAsync } from '../../../core/runtime'; import type { ImageBuffer } from '../image'; @@ -41,12 +41,14 @@ export async function createImageEmbedder( * Releases all allocated native resources. */ dispose: () => void; + /** * Asynchronously computes the embedding vector for the given input image. * @param input The input image buffer. * @returns A promise resolving to the embedding vector. */ embed: (input: ImageBuffer) => Promise; + /** * Synchronous version of {@link embed} to be executed directly on the * caller or worklet thread. @@ -56,14 +58,22 @@ export async function createImageEmbedder( const { modelPath, opts } = config; const model = await wrapAsync(loadModel, runtime)(modelPath); - const meta = validateModelSchema( - model, - 'forward', - [SymbolicTensor('float32', [1, 3, 'H', 'W'], [3, 'H', 'W'])], - [SymbolicTensor('float32', [1, 'D'], ['D'])] - ); - const inpShape = meta.inputTensorMeta[0]!.shape; - const outShape = meta.outputTensorMeta[0]!.shape; + const { variant, dims } = validateSpec(model.schema, { + batched: method( + 'forward', // prettier-ignore + [f32(1, 3, 'H', 'W')], + [f32(1, 'D')] + ), + unbatched: method( + 'forward', // prettier-ignore + [f32(3, 'H', 'W')], + [f32('D')] + ), + }); + + const [D, H, W] = dims.constant('D', 'H', 'W'); + const inpShape = { batched: [1, 3, H, W], unbatched: [3, H, W] }[variant]; + const outShape = { batched: [1, D], unbatched: [D] }[variant]; const tensors = [tensor('float32', outShape)] as const; const [tEmbedding] = tensors; diff --git a/packages/react-native-executorch/src/extensions/cv/tasks/instanceSegmentation.ts b/packages/react-native-executorch/src/extensions/cv/tasks/instanceSegmentation.ts index d26407ea39..59cdbc56d6 100644 --- a/packages/react-native-executorch/src/extensions/cv/tasks/instanceSegmentation.ts +++ b/packages/react-native-executorch/src/extensions/cv/tasks/instanceSegmentation.ts @@ -2,7 +2,7 @@ import type { WorkletRuntime } from 'react-native-worklets'; import { tensor } from '../../../core/tensor'; import { loadModel } from '../../../core/model'; -import { validateModelSchema, SymbolicTensor } from '../../../core/modelSchema'; +import { validateSpec, method, f32 } from '../../../core/schema'; import { wrapAsync } from '../../../core/runtime'; import type { ImageBuffer } from '../image'; @@ -115,35 +115,29 @@ export async function createInstanceSegmenter( }> { const { modelPath, opts } = config; const model = await wrapAsync(loadModel, runtime)(modelPath); - const meta = validateModelSchema( - model, - 'forward', - [SymbolicTensor('float32', [1, 3, 'H', 'W'], [3, 'H', 'W'])], - [ - SymbolicTensor('float32', ['N', 4]), - SymbolicTensor('float32', ['N']), - SymbolicTensor('float32', ['N']), - SymbolicTensor('float32', ['N', 'MH', 'MW']), - ] - ); - - const inpShape = meta.inputTensorMeta[0]!.shape; - - const outBoxesShape = meta.outputTensorMeta[0]!.shape; - const outScoresShape = meta.outputTensorMeta[1]!.shape; - const outClassesShape = meta.outputTensorMeta[2]!.shape; - const outMasksShape = meta.outputTensorMeta[3]!.shape; - - const maskH = outMasksShape[1]!; - const maskW = outMasksShape[2]!; - const targetH = inpShape.at(-2)!; - const targetW = inpShape.at(-1)!; + + const { variant, dims } = validateSpec(model.schema, { + batched: method( + 'forward', + [f32(1, 3, 'H', 'W')], + [f32('N', 4), f32('N'), f32('N'), f32('N', 'MH', 'MW')] + ), + unbatched: method( + 'forward', + [f32(3, 'H', 'W')], + [f32('N', 4), f32('N'), f32('N'), f32('N', 'MH', 'MW')] + ), + }); + + const [N, H, W, maskH, maskW] = dims.constant('N', 'H', 'W', 'MH', 'MW'); + const inpShape = { batched: [1, 3, H, W], unbatched: [3, H, W] }[variant]; + const outShape = { boxes: [N, 4], scores: [N], classes: [N], masks: [N, maskH, maskW] }; const tensors = [ - tensor('float32', outBoxesShape), - tensor('float32', outScoresShape), - tensor('float32', outClassesShape), - tensor('float32', outMasksShape), + tensor('float32', outShape.boxes), + tensor('float32', outShape.scores), + tensor('float32', outShape.classes), + tensor('float32', outShape.masks), tensor('float32', [maskH, maskW, 1]), ] as const; @@ -214,7 +208,7 @@ export async function createInstanceSegmenter( const d = boxes[idx * 4 + 3]!; const box = scaleBox(decodeBox([a, b, c, d], opts.boxFormat), { - from: { width: targetW, height: targetH }, + from: { width: W, height: H }, to: { width: input.width, height: input.height }, resizeMode: 'stretch', }); diff --git a/packages/react-native-executorch/src/extensions/cv/tasks/keypointDetection.ts b/packages/react-native-executorch/src/extensions/cv/tasks/keypointDetection.ts index d6ea03a294..e9c2e54085 100644 --- a/packages/react-native-executorch/src/extensions/cv/tasks/keypointDetection.ts +++ b/packages/react-native-executorch/src/extensions/cv/tasks/keypointDetection.ts @@ -2,7 +2,7 @@ import type { WorkletRuntime } from 'react-native-worklets'; import { tensor, type Tensor } from '../../../core/tensor'; import { loadModel } from '../../../core/model'; -import { validateModelSchema, SymbolicTensor } from '../../../core/modelSchema'; +import { validateSpec, method, f32 } from '../../../core/schema'; import { wrapAsync } from '../../../core/runtime'; import type { ImageBuffer } from '../image'; @@ -154,6 +154,7 @@ export async function createKeypointDetector void; + /** * Performs asynchronous keypoint and bounding box detection on the given * input image. @@ -169,6 +170,7 @@ export async function createKeypointDetector Promise[]>; + /** * Synchronous version of {@link detectKeypoints} to be executed directly on * the caller or worklet thread. @@ -181,27 +183,23 @@ export async function createKeypointDetector( * Releases all allocated native resources. */ dispose: () => void; + /** * Performs asynchronous object detection on the given input image. * @param input The input image buffer. @@ -84,6 +85,7 @@ export async function createObjectDetector( input: ImageBuffer, options?: { confidenceThreshold?: number; iouThreshold?: number } ) => Promise[]>; + /** * Synchronous version of {@link detectObjects} to be executed directly on the * caller or worklet thread. @@ -96,29 +98,27 @@ export async function createObjectDetector( const { modelPath, opts } = config; const model = await wrapAsync(loadModel, runtime)(modelPath); - const meta = validateModelSchema( - model, - 'forward', - [SymbolicTensor('float32', [1, 3, 'H', 'W'], [3, 'H', 'W'])], - [ - SymbolicTensor('float32', ['N', 4]), - SymbolicTensor('float32', ['N']), - SymbolicTensor('float32', ['N']), - ] - ); - - const inpShape = meta.inputTensorMeta[0]!.shape; - const outBoxesShape = meta.outputTensorMeta[0]!.shape; - const outScoresShape = meta.outputTensorMeta[1]!.shape; - const outClassesShape = meta.outputTensorMeta[2]!.shape; - - const targetH = inpShape.at(-2)!; - const targetW = inpShape.at(-1)!; + const { variant, dims } = validateSpec(model.schema, { + batched: method( + 'forward', // prettier-ignore + [f32(1, 3, 'H', 'W')], + [f32('N', 4), f32('N'), f32('N')] + ), + unbatched: method( + 'forward', // prettier-ignore + [f32(3, 'H', 'W')], + [f32('N', 4), f32('N'), f32('N')] + ), + }); + + const [N, H, W] = dims.constant('N', 'H', 'W'); + const inpShape = { batched: [1, 3, H, W], unbatched: [3, H, W] }[variant]; + const outShape = { boxes: [N, 4], scores: [N], classes: [N] }; const tensors = [ - tensor('float32', outBoxesShape), - tensor('float32', outScoresShape), - tensor('float32', outClassesShape), + tensor('float32', outShape.boxes), + tensor('float32', outShape.scores), + tensor('float32', outShape.classes), ] as const; const [tBoxes, tScores, tClasses] = tensors; @@ -176,7 +176,7 @@ export async function createObjectDetector( label, confidence, box: scaleBox(decodeBox([a, b, c, d], boxFormat), { - from: { width: targetW, height: targetH }, + from: { width: W, height: H }, to: { width: input.width, height: input.height }, ...opts, }), diff --git a/packages/react-native-executorch/src/extensions/cv/tasks/preprocessing.ts b/packages/react-native-executorch/src/extensions/cv/tasks/preprocessing.ts index 3569210481..9d0efd20bb 100644 --- a/packages/react-native-executorch/src/extensions/cv/tasks/preprocessing.ts +++ b/packages/react-native-executorch/src/extensions/cv/tasks/preprocessing.ts @@ -1,5 +1,4 @@ import { tensor, type Tensor } from '../../../core/tensor'; -import { matchShape } from '../../../core/modelSchema'; import type { ImageBuffer } from '../image'; import { @@ -62,15 +61,11 @@ export function createImagePreprocessor( dispose: () => void; } { const numRgbChannels = 3; - const expectedShapes = [ - [numRgbChannels, 'H', 'W'], - [1, numRgbChannels, 'H', 'W'], - ] as const; - - if (!matchShape(outputShape, ...expectedShapes)) { + const isRank3 = outputShape.length === 3 && outputShape[0] === numRgbChannels; + const isRank4 = outputShape.length === 4 && outputShape[1] === numRgbChannels; + if (!isRank3 && !isRank4) { throw new Error( - `preprocessor: got shape [${outputShape}], required one of: ` + - `${expectedShapes.map((s) => `[${s.join(',')}]`).join(' | ')}` + `preprocessor: got shape [${outputShape}], expected [${numRgbChannels}, H, W] or [1, ${numRgbChannels}, H, W]` ); } diff --git a/packages/react-native-executorch/src/extensions/cv/tasks/sdxsTextToImage.ts b/packages/react-native-executorch/src/extensions/cv/tasks/sdxsTextToImage.ts index 7747a4517c..0130e6cd07 100644 --- a/packages/react-native-executorch/src/extensions/cv/tasks/sdxsTextToImage.ts +++ b/packages/react-native-executorch/src/extensions/cv/tasks/sdxsTextToImage.ts @@ -2,7 +2,7 @@ import type { WorkletRuntime } from 'react-native-worklets'; import { tensor } from '../../../core/tensor'; import { loadModel } from '../../../core/model'; -import { validateModelSchema, SymbolicTensor } from '../../../core/modelSchema'; +import { validateSpec, method, i64, f32 } from '../../../core/schema'; import { wrapAsync } from '../../../core/runtime'; import { randomNormal } from '../../math'; import { loadTokenizer } from '../../nlp/tokenizer'; @@ -21,10 +21,9 @@ const CLIP_PAD_TOKEN_ID = 49407; const IMAGE_SIZE = 512; const LATENT_CHANNELS = 4; -// The UNet operates at 1/8 of the image resolution. -const LATENT_SIZE = IMAGE_SIZE / 8; -// Scheduler `init_noise_sigma`: the initial latents are standard normal. -const INIT_NOISE_SIGMA = 1.0; +const LATENT_SIZE = IMAGE_SIZE / 8; // The UNet operates at 1/8 of the image resolution. +const LATENT_SHAPE = [1, LATENT_CHANNELS, LATENT_SIZE, LATENT_SIZE]; +const INIT_NOISE_SIGMA = 1.0; // Scheduler's initial latents are standard normal. // At the distilled single step the DEIS update collapses to an exactly linear // combination of the latents and the UNet output: @@ -61,6 +60,7 @@ export async function createSdxsTextToImage( ): Promise<{ /** Releases all allocated native resources. */ dispose: () => void; + /** * Generates an image from a text prompt. * @param prompt The text prompt describing the desired image. @@ -69,6 +69,7 @@ export async function createSdxsTextToImage( * @returns A promise resolving to the generated RGBA image buffer. */ generate: (prompt: string, seed?: number) => Promise; + /** * Synchronous version of {@link generate} to be executed directly on the * caller or worklet thread. @@ -79,35 +80,32 @@ export async function createSdxsTextToImage( const model = await wrapAsync(loadModel, runtime)(modelPath); const tokenizer = await wrapAsync(loadTokenizer, runtime)(tokenizerPath); - validateModelSchema( - model, - 'encode', - [SymbolicTensor('int64', [1, CLIP_MAX_TOKENS])], - [SymbolicTensor('float32', [1, CLIP_MAX_TOKENS, CLIP_HIDDEN_SIZE])] - ); - validateModelSchema( - model, - 'denoise', - [ - SymbolicTensor('float32', [1, LATENT_CHANNELS, LATENT_SIZE, LATENT_SIZE]), - SymbolicTensor('int64', [1]), - SymbolicTensor('float32', [1, CLIP_MAX_TOKENS, CLIP_HIDDEN_SIZE]), - ], - [SymbolicTensor('float32', [1, LATENT_CHANNELS, LATENT_SIZE, LATENT_SIZE])] - ); - validateModelSchema( - model, - 'decode', - [SymbolicTensor('float32', [1, LATENT_CHANNELS, LATENT_SIZE, LATENT_SIZE])], - [SymbolicTensor('float32', [1, 3, IMAGE_SIZE, IMAGE_SIZE])] - ); + validateSpec(model.schema, { + default: { + ...method( + 'encode', // prettier-ignore + [i64(1, CLIP_MAX_TOKENS)], + [f32(1, CLIP_MAX_TOKENS, CLIP_HIDDEN_SIZE)] + ), + ...method( + 'denoise', + [f32(...LATENT_SHAPE), i64(1), f32(1, CLIP_MAX_TOKENS, CLIP_HIDDEN_SIZE)], + [f32(...LATENT_SHAPE)] + ), + ...method( + 'decode', // prettier-ignore + [f32(...LATENT_SHAPE)], + [f32(1, 3, IMAGE_SIZE, IMAGE_SIZE)] + ), + }, + }); const tensors = [ tensor('int64', [1, CLIP_MAX_TOKENS]), tensor('float32', [1, CLIP_MAX_TOKENS, CLIP_HIDDEN_SIZE]), tensor('int64', [1]), - tensor('float32', [1, LATENT_CHANNELS, LATENT_SIZE, LATENT_SIZE]), - tensor('float32', [1, LATENT_CHANNELS, LATENT_SIZE, LATENT_SIZE]), + tensor('float32', LATENT_SHAPE), + tensor('float32', LATENT_SHAPE), tensor('float32', [1, 3, IMAGE_SIZE, IMAGE_SIZE]), tensor('float32', [3, IMAGE_SIZE, IMAGE_SIZE]), tensor('uint8', [3, IMAGE_SIZE, IMAGE_SIZE]), diff --git a/packages/react-native-executorch/src/extensions/cv/tasks/semanticSegmentation.ts b/packages/react-native-executorch/src/extensions/cv/tasks/semanticSegmentation.ts index 148cf9b267..b1b19b0385 100644 --- a/packages/react-native-executorch/src/extensions/cv/tasks/semanticSegmentation.ts +++ b/packages/react-native-executorch/src/extensions/cv/tasks/semanticSegmentation.ts @@ -2,7 +2,7 @@ import type { WorkletRuntime } from 'react-native-worklets'; import { tensor } from '../../../core/tensor'; import { loadModel } from '../../../core/model'; -import { validateModelSchema, SymbolicTensor } from '../../../core/modelSchema'; +import { validateSpec, method, f32 } from '../../../core/schema'; import { wrapAsync } from '../../../core/runtime'; import type { ImageBuffer } from '../image'; @@ -84,6 +84,7 @@ export async function createSemanticSegmenter( * Releases all allocated native resources. */ dispose: () => void; + /** * Runs semantic segmentation asynchronously. * @@ -109,6 +110,7 @@ export async function createSemanticSegmenter( input: ImageBuffer, colormap?: Partial> ) => Promise>; + /** * Runs semantic segmentation synchronously. * @see {@link segment} for details. @@ -121,18 +123,22 @@ export async function createSemanticSegmenter( const { modelPath, opts } = config; const model = await wrapAsync(loadModel, runtime)(modelPath); - const meta = validateModelSchema( - model, - 'forward', - [SymbolicTensor('float32', [1, 3, 'H', 'W'], [3, 'H', 'W'])], - [SymbolicTensor('float32', [1, 'K', 'H', 'W'], ['K', 'H', 'W'])] - ); - const inpShape = meta.inputTensorMeta[0]!.shape; - const outShape = meta.outputTensorMeta[0]!.shape; + const { variant, dims } = validateSpec(model.schema, { + batched: method( + 'forward', // prettier-ignore + [f32(1, 3, 'H', 'W')], + [f32(1, 'K', 'H', 'W')] + ), + unbatched: method( + 'forward', // prettier-ignore + [f32(3, 'H', 'W')], + [f32('K', 'H', 'W')] + ), + }); - const nClasses = outShape.at(-3)!; - const targetH = outShape.at(-2)!; - const targetW = outShape.at(-1)!; + const [nClasses, H, W] = dims.constant('K', 'H', 'W'); + const inpShape = { batched: [1, 3, H, W], unbatched: [3, H, W] }[variant]; + const outShape = { batched: [1, nClasses, H, W], unbatched: [nClasses, H, W] }[variant]; // Generate highly distinct, high-contrast colors, see: // https://martin.ankerl.com/2009/12/09/how-to-create-random-colors-programmatically/ @@ -149,11 +155,11 @@ export async function createSemanticSegmenter( const tensors = [ tensor('float32', outShape), - tensor('float32', [nClasses, targetH, targetW]), - tensor('float32', [nClasses, targetH, targetW]), - tensor('float32', [targetH, targetW, nClasses]), - tensor(nClasses > 1 ? 'int32' : 'uint8', [targetH, targetW, 1]), - tensor('uint8', [targetH, targetW, 4]), + tensor('float32', [nClasses, H, W]), + tensor('float32', [nClasses, H, W]), + tensor('float32', [H, W, nClasses]), + tensor(nClasses > 1 ? 'int32' : 'uint8', [H, W, 1]), + tensor('uint8', [H, W, 4]), ] as const; const [tOutput, tReshape, tSigmoid, tChanLast, tMask, tRgba] = tensors; diff --git a/packages/react-native-executorch/src/extensions/cv/tasks/styleTransfer.ts b/packages/react-native-executorch/src/extensions/cv/tasks/styleTransfer.ts index 1c1a89cbe5..ede6b88f0d 100644 --- a/packages/react-native-executorch/src/extensions/cv/tasks/styleTransfer.ts +++ b/packages/react-native-executorch/src/extensions/cv/tasks/styleTransfer.ts @@ -2,7 +2,7 @@ import type { WorkletRuntime } from 'react-native-worklets'; import { tensor } from '../../../core/tensor'; import { loadModel } from '../../../core/model'; -import { validateModelSchema, SymbolicTensor } from '../../../core/modelSchema'; +import { validateSpec, method, f32 } from '../../../core/schema'; import { wrapAsync } from '../../../core/runtime'; import type { ImageBuffer } from '../image'; @@ -54,12 +54,14 @@ export async function createStyleTransfer( * Releases all allocated native resources. */ dispose: () => void; + /** * Performs asynchronous image style transfer on the given input image. * @param input The input image buffer. * @returns A promise resolving to the styled image buffer. */ transferStyle: (input: ImageBuffer) => Promise; + /** * Synchronous version of {@link transferStyle} to be executed directly on the * caller or worklet thread. @@ -69,24 +71,29 @@ export async function createStyleTransfer( const { modelPath, opts } = config; const model = await wrapAsync(loadModel, runtime)(modelPath); - const meta = validateModelSchema( - model, - 'forward', - [SymbolicTensor('float32', [1, 3, 'H', 'W'], [3, 'H', 'W'])], - [SymbolicTensor('float32', [1, 3, 'H', 'W'], [3, 'H', 'W'])] - ); - const inpShape = meta.inputTensorMeta[0]!.shape; - const outShape = meta.outputTensorMeta[0]!.shape; + const { variant, dims } = validateSpec(model.schema, { + batched: method( + 'forward', // prettier-ignore + [f32(1, 3, 'H', 'W')], + [f32(1, 3, 'H', 'W')] + ), + unbatched: method( + 'forward', // prettier-ignore + [f32(3, 'H', 'W')], + [f32(3, 'H', 'W')] + ), + }); - const targetH = outShape.at(-2)!; - const targetW = outShape.at(-1)!; + const [H, W] = dims.constant('H', 'W'); + const inpShape = { batched: [1, 3, H, W], unbatched: [3, H, W] }[variant]; + const outShape = inpShape; const tensors = [ tensor('float32', outShape), - tensor('float32', [3, targetH, targetW]), - tensor('uint8', [3, targetH, targetW]), - tensor('uint8', [targetH, targetW, 3]), - tensor('uint8', [targetH, targetW, 4]), + tensor('float32', [3, H, W]), + tensor('uint8', [3, H, W]), + tensor('uint8', [H, W, 3]), + tensor('uint8', [H, W, 4]), ] as const; const [tOutput, tReshape, tUint8, tChanLast, tRgba] = tensors; diff --git a/packages/react-native-executorch/src/extensions/nlp/tasks/textEmbedding.ts b/packages/react-native-executorch/src/extensions/nlp/tasks/textEmbedding.ts index 857457d89f..3aa88f0b05 100644 --- a/packages/react-native-executorch/src/extensions/nlp/tasks/textEmbedding.ts +++ b/packages/react-native-executorch/src/extensions/nlp/tasks/textEmbedding.ts @@ -2,7 +2,7 @@ import type { WorkletRuntime } from 'react-native-worklets'; import { tensor } from '../../../core/tensor'; import { loadModel } from '../../../core/model'; -import { validateModelSchema, SymbolicTensor } from '../../../core/modelSchema'; +import { validateSpec, DynamicDim as Dyn, method, i64, f32, constr } from '../../../core/schema'; import { wrapAsync } from '../../../core/runtime'; import { loadTokenizer } from '../tokenizer'; @@ -44,6 +44,7 @@ export async function createTextEmbedder( * Releases all allocated native resources. */ dispose: () => void; + /** * Asynchronously computes the embedding vector for the given input text. * Inputs longer than the model's maximum sequence length are truncated. @@ -53,6 +54,7 @@ export async function createTextEmbedder( * @returns A promise resolving to the embedding vector. */ embed: (input: string, prompt?: string) => Promise; + /** * Synchronous version of {@link embed} to be executed directly on the * caller or worklet thread. @@ -67,16 +69,34 @@ export async function createTextEmbedder( // Text embedding models take two int64 inputs: the token ids and the // attention mask, both of shape [1, sequence_length]. - const meta = validateModelSchema( - model, - 'forward', - [SymbolicTensor('int64', [1, 'L']), SymbolicTensor('int64', [1, 'L'])], - [SymbolicTensor('float32', [1, 'D'], ['D'])] - ); - // The models are exported with a dynamic sequence dimension; the declared size - // is the upper bound, used only to truncate over-long inputs. - const maxSeqLen = meta.inputTensorMeta[0]!.shape[1]!; - const outShape = meta.outputTensorMeta[0]!.shape; + const { variant, dims } = validateSpec(model.schema, { + batched: method( + 'forward', // prettier-ignore + [i64(1, Dyn('L')), i64(1, Dyn('L'))], + [f32(1, 'D')], + [ + constr.eq( + { paramSide: 'input', tensorIdx: 0, dimIdx: 1 }, + { paramSide: 'input', tensorIdx: 1, dimIdx: 1 } + ), + ] + ), + unbatched: method( + 'forward', // prettier-ignore + [i64(1, Dyn('L')), i64(1, Dyn('L'))], + [f32('D')], + [ + constr.eq( + { paramSide: 'input', tensorIdx: 0, dimIdx: 1 }, + { paramSide: 'input', tensorIdx: 1, dimIdx: 1 } + ), + ] + ), + }); + + const [seqLen] = dims.range('L'); + const [D] = dims.constant('D'); + const outShape = { batched: [1, D], unbatched: [D] }[variant]; const tensors = [tensor('float32', outShape)] as const; const [tEmbedding] = tensors; @@ -94,7 +114,7 @@ export async function createTextEmbedder( if (ids.length === 0) { throw new Error('createTextEmbedder: input tokenized to zero tokens'); } - const len = Math.min(ids.length, maxSeqLen); + const len = Math.min(ids.length, seqLen.max); const idsData = new BigInt64Array(len); const maskData = new BigInt64Array(len); diff --git a/packages/react-native-executorch/src/extensions/speech/tasks/fsmnVoiceActivityDetection.ts b/packages/react-native-executorch/src/extensions/speech/tasks/fsmnVoiceActivityDetection.ts index 9225c743ff..0fd6f4885a 100644 --- a/packages/react-native-executorch/src/extensions/speech/tasks/fsmnVoiceActivityDetection.ts +++ b/packages/react-native-executorch/src/extensions/speech/tasks/fsmnVoiceActivityDetection.ts @@ -2,7 +2,8 @@ import type { WorkletRuntime } from 'react-native-worklets'; import { tensor } from '../../../core/tensor'; import { loadModel } from '../../../core/model'; -import { validateModelSchema, SymbolicTensor } from '../../../core/modelSchema'; +import { validateSpec, DynamicDim as Dyn, method, f32 } from '../../../core/schema'; + import { wrapAsync } from '../../../core/runtime'; import { extractFrames } from '../utils/vadUtils'; @@ -187,6 +188,7 @@ export async function createFsmnVoiceActivityDetector( * Releases all allocated native resources. */ dispose: () => void; + /** * Asynchronously detects speech segments within a mono waveform sampled at * {@link FSMN_VAD_SAMPLE_RATE_HZ}. @@ -195,11 +197,13 @@ export async function createFsmnVoiceActivityDetector( * @returns A promise resolving to the detected speech segments, in seconds. */ detectVoice: (waveform: Float32Array, options?: VadOptions) => Promise; + /** * Synchronous version of {@link detectVoice} to be executed directly on the caller * or worklet thread. */ detectVoiceWorklet: (waveform: Float32Array, options?: VadOptions) => Segment[]; + /** * Appends a live audio chunk to a bounded rolling window, runs detection over * that window and reports a {@link VadEvent} when speech starts or stops, @@ -209,6 +213,7 @@ export async function createFsmnVoiceActivityDetector( * @param options Optional overrides of the detection thresholds and margin. */ detectVoiceOnStream: (chunk: Float32Array, options?: VadStreamOptions) => VadEvent | undefined; + /** * Clears the rolling window and speaking state used by {@link detectVoiceOnStream}. */ @@ -222,15 +227,15 @@ export async function createFsmnVoiceActivityDetector( // [1, frames, classes] where class 0 is the non-speech class. The output frame // count matches the input at runtime, but ExecuTorch metadata only reports the // static upper bound, so the per-call output tensor is sized explicitly below. - const meta = validateModelSchema( - model, - 'forward', - [SymbolicTensor('float32', ['frames', 'fftLength'])], - [SymbolicTensor('float32', [1, 'frames', 'classes'])] - ); - const maxFrames = meta.inputTensorMeta[0]!.shape[0]!; - const fftLength = meta.inputTensorMeta[0]!.shape[1]!; - const numClass = meta.outputTensorMeta[0]!.shape[2]!; + const { dims } = validateSpec(model.schema, { + default: method( + 'forward', // prettier-ignore + [f32(Dyn('frames'), 'fftLen')], + [f32(1, Dyn('frames'), 'classes')] + ), + }); + const [numClasses, fftLength] = dims.constant('classes', 'fftLen'); + const maxFrames = dims.range('frames')[0].max; // The Hann window is uploaded once and reused by the native framing op across // every call. @@ -264,7 +269,7 @@ export async function createFsmnVoiceActivityDetector( waveform.subarray(startSample, startSample + sampleCount) ); const tInput = tensor('float32', [chunkFrames, fftLength]); - const tOutput = tensor('float32', [1, chunkFrames, numClass]); + const tOutput = tensor('float32', [1, chunkFrames, numClasses]); const outBuffer = new Float32Array(tOutput.numel); try { extractFrames(tWaveform, tHann, tInput, { @@ -281,7 +286,7 @@ export async function createFsmnVoiceActivityDetector( } for (let i = 0; i < realFrames; i++) { - scores[offset + i] = outBuffer[i * numClass]!; + scores[offset + i] = outBuffer[i * numClasses]!; } offset += realFrames; } diff --git a/packages/react-native-executorch/src/extensions/speech/tasks/whisperSpeechToText.ts b/packages/react-native-executorch/src/extensions/speech/tasks/whisperSpeechToText.ts index dcd71f9bc4..d5a202eaca 100644 --- a/packages/react-native-executorch/src/extensions/speech/tasks/whisperSpeechToText.ts +++ b/packages/react-native-executorch/src/extensions/speech/tasks/whisperSpeechToText.ts @@ -3,8 +3,8 @@ import { scheduleOnRN, createSynchronizable } from 'react-native-worklets'; import { tensor } from '../../../core/tensor'; import { loadModel } from '../../../core/model'; -import { validateModelSchema, SymbolicTensor } from '../../../core/modelSchema'; import { wrapAsync } from '../../../core/runtime'; +import { f32, i64, method, validateSpec, DynamicDim as Dyn } from '../../../core/schema'; import { argmax } from '../../../extensions/math'; import { loadTokenizer } from '../../nlp/tokenizer'; @@ -166,32 +166,28 @@ export async function createWhisperSpeechToText')!; const isEnglishOnly = supportedLanguages.length === 1 && supportedLanguages[0] === 'en'; - const encMeta = validateModelSchema( - model, - 'encode', - [SymbolicTensor('float32', ['T_audio'])], - [SymbolicTensor('float32', [1, 'Seq', 'State'])] - ); - - const encSeqLen = encMeta.outputTensorMeta[0]!.shape[1]!; - const encStateDim = encMeta.outputTensorMeta[0]!.shape[2]!; - - validateModelSchema( - model, - 'decode', - [ - SymbolicTensor('int64', [1, 'Tokens']), - SymbolicTensor('int64', ['Tokens']), - SymbolicTensor('float32', [1, encSeqLen, encStateDim]), - ], - [SymbolicTensor('float32', [1, 'Tokens', 'Vocab'])] - ); + const { dims } = validateSpec(model.schema, { + default: { + ...method( + 'encode', // prettier-ignore + [f32(Dyn('T_audio'))], // Internally, model always pads audio to the maximum duration. + [f32(1, 'SeqLen', 'StateDim')] + ), + ...method( + 'decode', + [i64(1, 1), i64(1), f32(1, 'SeqLen', 'StateDim')], + [f32(1, 1, 'VocabSize')] + ), + }, + }); + + const [seqLen, stateDim] = dims.constant('SeqLen', 'StateDim'); const tensors = [ tensor('int64', [1]), // tPosition tensor('int64', [1, 1]), // tToken tensor('int32', [1, 1, 1]), // tArgmax - tensor('float32', [1, encSeqLen, encStateDim]), //tEncodings + tensor('float32', [1, seqLen, stateDim]), //tEncodings tensor('float32', [1, 1, tokenizer.getVocabSize()]), // tLogits ] as const; diff --git a/packages/react-native-executorch/src/index.ts b/packages/react-native-executorch/src/index.ts index 3aad5c8d77..1e429d0adf 100644 --- a/packages/react-native-executorch/src/index.ts +++ b/packages/react-native-executorch/src/index.ts @@ -33,23 +33,12 @@ export * from './extensions/speech/tasks/fsmnVoiceActivityDetection'; export * from './extensions/speech/tasks/whisperSpeechToText'; // Core primitives — for library builders and power users -export { tensor } from './core/tensor'; -export type { DType, Tensor } from './core/tensor'; +export * from './core/model'; +export * from './core/tensor'; +export * from './core/runtime'; -export { loadModel } from './core/model'; -export type { - Model, - ModelInput, - ModelOutput, - TensorMeta, - ModelMethodMeta, - ExecuTorchTag, -} from './core/model'; - -export { validateModelSchema, SymbolicTensor, matchShape } from './core/modelSchema'; -export type { ValueConstraint, TensorConstraint, SymbolicShape } from './core/modelSchema'; - -export { defaultWorkletRuntime, wrapAsync } from './core/runtime'; +export type * from './core/schema'; +export * as schema from './core/schema'; export * as math from './extensions/math'; export * as cv from './extensions/cv'; diff --git a/packages/react-native-executorch/src/utils.ts b/packages/react-native-executorch/src/utils.ts index 1c457b3062..154f7c2aa4 100644 --- a/packages/react-native-executorch/src/utils.ts +++ b/packages/react-native-executorch/src/utils.ts @@ -1,5 +1,6 @@ import { rnexecutorchJsi } from './native/bridge'; -import { loadModel, type ModelMethodMeta } from './core/model'; +import { loadModel } from './core/model'; +import type { ModelSpec, ConcreteDim } from './core/schema'; import RNFS from 'react-native-fs'; /** @@ -24,12 +25,13 @@ export function getRegisteredBackends(): string[] { * @category Utils * @experimental Subject to change once the temporary react-native-fs dependency is replaced. See [Issue #1253](https://github.com/software-mansion/react-native-executorch/issues/1253). * @param source The remote HTTP URL or local path to the `.pte` model file. - * @returns A promise resolving to an object containing the model source and - * method signature metadata. + * @returns A promise resolving to an object containing the model source, + * method signature metadata, and per-method backend usage. */ export async function inspectModel(source: string): Promise<{ source: string; - methods: { name: string; meta: ModelMethodMeta }[]; + schema: ModelSpec; + backends: Record; }> { let localPath = source; let downloaded = false; @@ -44,15 +46,7 @@ export async function inspectModel(source: string): Promise<{ try { model = loadModel(localPath); - const methodNames = model.getMethodNames(); - - const methods: { name: string; meta: ModelMethodMeta }[] = []; - for (const method of methodNames) { - const meta = model.getMethodMeta(method); - methods.push({ name: method, meta }); - } - - return { source, methods }; + return { source, schema: model.schema, backends: model.backends }; } finally { if (model) { model.dispose();