Add EmbeddingGemma 2 export: coreai.embedding.export and registry entry

models/README.md

@@ -133,7 +133,7 @@ Pass `--include-debug-info` to switch the converter to `DEBUG` mode, which embed
 uv run coreai.llm.export Qwen/Qwen3-0.6B --include-debug-info
 ```
 
-The flag is available on `coreai.llm.export`, `coreai.vlm.export`, `coreai.diffusion.export`, `coreai.segmentation.export`, and every standalone `models/<name>/export.py` recipe, so all export paths produce assets carrying the same debug information by default.
+The flag is available on `coreai.llm.export`, `coreai.vlm.export`, `coreai.diffusion.export`, `coreai.segmentation.export`, `coreai.embedding.export`, and every standalone `models/<name>/export.py` recipe, so all export paths produce assets carrying the same debug information by default.
 
 **Note:** `--include-debug-info` is independent of `--verbose`/`-v`. `--verbose` only raises the console log level; it does not change what goes into the asset.
 
@@ -189,6 +189,16 @@ decoder (`main`), token-embedding lookup (`embedding`), vision encoder
 (`vision`), tokenizer, and `metadata.json`. Pass `--skip-vision` to export the
 text portion only.
 
+### Embedding Models
+
+```bash
+uv run coreai.embedding.export --list-models
+uv run coreai.embedding.export embeddinggemma-2                 # text + image (macOS)
+uv run coreai.embedding.export embeddinggemma-2 --platform iOS  # Neural Engine
+```
+
+See the [EmbeddingGemma 2](embedding_gemma2) model card.
+
 ### Standalone Export Scripts
 
 Models with a standalone `export.py` are run directly:
@@ -245,9 +255,13 @@ uv run models/<name>/export.py --include-debug-info   # embed debug information
 - [RoBERTa](roberta)
 - [T5](t5)
 
+### Embedding Models
+
+- [EmbeddingGemma 2](embedding_gemma2)
+
 ## Adding a Model
 
-To make a new model exportable via short-name, add a `ModelPreset(...)` entry to `LLM_PRESETS` or `DIFFUSION_PRESETS` in [`python/src/coreai_models/model_registry.py`](../python/src/coreai_models/model_registry.py). Set the short name, HuggingFace ID, family, variant, and the export defaults (compression, compute precision, max context length).
+To make a new model exportable via short-name, add a `ModelPreset(...)` entry to `LLM_PRESETS` or `DIFFUSION_PRESETS` in [`python/src/coreai_models/model_registry.py`](../python/src/coreai_models/model_registry.py). Set the short name, HuggingFace ID, family, variant, and the export defaults (compression, compute precision, max context length). Models exported by `coreai.embedding.export` need an `EmbeddingFamily` in `EMBEDDING_FAMILIES` ([`python/src/coreai_models/embedding/contract.py`](../python/src/coreai_models/embedding/contract.py)) and a `UtilityModel` entry with task `embedding` in `UTILITY_PRESETS`.
 
 For models with bespoke export logic that doesn't fit the standard `coreai.llm.export` / `coreai.diffusion.export` flow, write a standalone recipe under `models/<name>/export.py` — see existing recipes for the [PEP 723](https://peps.python.org/pep-0723/) pattern and `models/README.md` for the contribution checklist.
 

models/embedding_gemma2/README.md

@@ -0,0 +1,57 @@
+# EmbeddingGemma 2
+
+Multimodal retrieval embedding model that maps text and images into one shared space: each input becomes a single 768-dimensional, L2-normalized vector, and its leading 512, 256 or 128 dimensions, re-normalized, are valid smaller embeddings.[^1]
+
+## Setup
+
+If you haven't installed `uv`, install it by
+
+```bash
+brew install uv
+```
+
+## Export
+
+```sh
+uv run coreai.embedding.export embeddinggemma-2
+```
+
+Saves to `<repo-root>/exports/embeddinggemma-2/` (text and image towers, GPU). `--platform iOS` writes `embeddinggemma-2-ios/` and `--skip-vision` appends `-text`. Pass `--output-dir <path>` to change the parent directory.
+
+```sh
+uv run coreai.embedding.export --help
+```
+
+**Options:**
+
+| Flag                   | Description                                          | Default                              |
+| ---------------------- | ---------------------------------------------------- | ------------------------------------ |
+| `--platform`           | `macOS` (GPU bundle) or `iOS` (Neural Engine bundle) | `macOS`                              |
+| `--skip-vision`        | Export the text tower only                           | —                                    |
+| `--compute-precision`  | `float32`, `bfloat16` on macOS; `float16` on iOS     | `float32` on macOS, `float16` on iOS |
+| `--output-dir`         | Parent directory for the bundle                      | `<repo-root>/exports/`               |
+| `--overwrite`          | Overwrite an existing bundle                         | —                                    |
+| `--include-debug-info` | Embed conversion debug info in the `.aimodel` files  | —                                    |
+| `--dry-run`            | Print the export plan without writing                | —                                    |
+| `--list-models`        | List the supported models and exit                   | —                                    |
+
+**Supported models:**
+
+| Model            | HuggingFace ID          | Parameters                   |
+| ---------------- | ----------------------- | ---------------------------- |
+| embeddinggemma-2 | google/embeddinggemma-2 | 439M (text 271M, image 168M) |
+
+## Bundle
+
+`metadata.json` lists each `.aimodel` file once under `assets`. Each entry of `embedding.components` names the asset it lives in (`asset`), its functions and its `output_dtype`:
+
+| Component       | macOS asset (GPU)            | iOS asset                       | Functions          | Output dtype                   |
+| --------------- | ---------------------------- | ------------------------------- | ------------------ | ------------------------------ |
+| `text_encoder`  | `text_encoder`               | —                               | `len{L}_b{B}`      | `float32`                      |
+| `token_table`   | `text_encoder` (images only) | `token_table` (GPU)             | `main`             | `float32` macOS, `float16` iOS |
+| `encoder`       | `text_encoder` (images only) | `encoder` (Neural Engine)       | `embeds_len{L}_b1` | `float32` macOS, `float16` iOS |
+| `image_encoder` | `image_encoder`              | `image_encoder` (Neural Engine) | `main`             | `float32` macOS, `float16` iOS |
+
+The lengths `{L}` differ by platform; read the function names from `embedding.components.<component>.functions`.
+
+[^1]: [Documentation](https://ai.google.dev/gemma/docs/embeddinggemma) · [HuggingFace](https://huggingface.co/google/embeddinggemma-2)

python/pyproject.toml

@@ -49,6 +49,7 @@ dependencies = [
 "coreai.llm.export" = "coreai_models.llm.export:main"
 "coreai.llm.eval" = "coreai_models.llm.eval:main"
 "coreai.vlm.export" = "coreai_models.vlm.export:main"
+"coreai.embedding.export" = "coreai_models.embedding.export:main"
 "coreai.diffusion.export" = "coreai_models.diffusion.export:main"
 "coreai.model.registry" = "coreai_models.model_registry:main"
 

python/src/coreai_models/diffusion/gpu.py

@@ -82,13 +82,7 @@ def export_fn(module: torch.nn.Module) -> torch.export.ExportedProgram:
         decomposed: torch.export.ExportedProgram = exported.run_decompositions(coreai_decomp_table)
         return decomposed
 
-    converter = coreai_torch.TorchConverter(
-        mode=(
-            coreai_torch.TorchConverter.Mode.DEBUG
-            if include_debug_info
-            else coreai_torch.TorchConverter.Mode.RELEASE
-        )
-    )
+    converter = make_converter(include_debug_info)
     converter.add_pytorch_module(
         wrapper,
         export_fn=export_fn,
@@ -99,6 +93,42 @@ def export_fn(module: torch.nn.Module) -> torch.export.ExportedProgram:
     return program
 
 
+def make_converter(
+    include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO,
+) -> coreai_torch.TorchConverter:
+    """Return a ``TorchConverter``; ``include_debug_info`` keeps debug info in the asset."""
+    mode = coreai_torch.TorchConverter.Mode
+    return coreai_torch.TorchConverter(mode=mode.DEBUG if include_debug_info else mode.RELEASE)
+
+
+def add_function(
+    converter: coreai_torch.TorchConverter,
+    name: str,
+    module: torch.nn.Module,
+    args: tuple[torch.Tensor, ...],
+    input_names: tuple[str, ...] | list[str],
+    output_names: tuple[str, ...] | list[str],
+    *,
+    dynamic_shapes: tuple[dict[int, torch.export.Dim] | None, ...] | None = None,
+    decompositions: dict | None = None,
+) -> None:
+    """Trace ``module`` on ``args`` and add it to ``converter`` as function ``name``.
+
+    Functions added to one converter share their weights. ``decompositions`` defaults to Core AI's
+    table.
+    """
+    module.eval()
+    with torch.no_grad():
+        exported = torch.export.export(module, args=args, dynamic_shapes=dynamic_shapes)
+    table = coreai_torch.get_decomp_table() if decompositions is None else decompositions
+    converter.add_exported_program(
+        exported.run_decompositions(table),
+        input_names=input_names,
+        output_names=output_names,
+        entrypoint_name=name,
+    )
+
+
 def export_multifunction(
     functions: list[tuple[str, torch.nn.Module, tuple[torch.Tensor, ...]]],
     input_names: tuple[str, ...],
@@ -122,29 +152,8 @@ def export_multifunction(
     Returns:
         An optimized AIProgram with multiple named functions.
     """
-    converter = coreai_torch.TorchConverter(
-        mode=(
-            coreai_torch.TorchConverter.Mode.DEBUG
-            if include_debug_info
-            else coreai_torch.TorchConverter.Mode.RELEASE
-        )
-    )
-    coreai_decomp_table = coreai_torch.get_decomp_table()
-
+    converter = make_converter(include_debug_info)
     for name, wrapper, dummy_inputs in functions:
-        wrapper.eval()
         logger.info(f"Tracing function '{name}' (seq dims: {[t.shape for t in dummy_inputs[:2]]})")
-
-        with torch.no_grad():
-            exported = torch.export.export(wrapper, args=dummy_inputs)
-        decomposed = exported.run_decompositions(coreai_decomp_table)
-
-        converter.add_exported_program(
-            decomposed,
-            input_names=input_names,
-            output_names=output_names,
-            entrypoint_name=name,
-        )
-
-    program = converter.to_coreai()
-    return program
+        add_function(converter, name, wrapper, dummy_inputs, input_names, output_names)
+    return converter.to_coreai()

python/src/coreai_models/embedding/__init__.py

@@ -0,0 +1,6 @@
+# Copyright 2026 Apple Inc.
+#
+# Use of this source code is governed by a BSD-3-clause license that can
+# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
+
+"""Retrieval-embedding bundle export; CLI: `uv run coreai.embedding.export`."""