Export an MLX model for the iOS demo build

.ci/scripts/test_ios_ci.sh

@@ -49,6 +49,43 @@ python3 -m examples.portable.scripts.export --model_name="$MODEL_NAME" --segment
 python3 -m examples.apple.coreml.scripts.export --model_name="$MODEL_NAME"
 python3 -m examples.xnnpack.aot_compiler --model_name="$MODEL_NAME" --delegate
 
+# No generic examples/ entry point takes a model name and lowers it to MLX, so do
+# it here. Lowering is ahead of time and imports no mlx runtime package. Assert MLX
+# claimed the graph before writing: the partitioner only warns when it delegates
+# nothing, so without this a program with no MLX delegate would be staged.
+python3 - "$MODEL_NAME" <<'PY'
+import sys
+
+import torch
+from executorch.backends.mlx import MLXPartitioner
+from executorch.examples.models import MODEL_NAME_TO_MODEL
+from executorch.examples.models.model_factory import EagerModelFactory
+from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower
+
+model_name = sys.argv[1]
+model, example_inputs, _, _ = EagerModelFactory.create_model(
+ *MODEL_NAME_TO_MODEL[model_name]
+)
+program = to_edge_transform_and_lower(
+ torch.export.export(model.eval(), example_inputs),
+ partitioner=[MLXPartitioner()],
+ compile_config=EdgeCompileConfig(_skip_dim_order=True),
+).to_executorch()
+
+segments = sum(
+ delegate.id == "MLXBackend"
+ for plan in program.executorch_program.execution_plan
+ for delegate in plan.delegates
+)
+if segments == 0:
+ raise SystemExit(f"error: {model_name} produced no MLX delegate")
+
+path = f"{model_name}_mlx.pte"
+with open(path, "wb") as file:
+ program.write_to_file(file)
+print(f"{path}: {segments} MLX delegate segment(s)")
+PY
+
 mkdir -p "$APP_PATH/Resources/Models/MobileNet/"
 mv $MODEL_NAME*.pte "$APP_PATH/Resources/Models/MobileNet/"
 

scripts/test_ios.sh

@@ -66,6 +66,43 @@ python3 -m examples.portable.scripts.export --model_name="$MODEL_NAME"
 python3 -m examples.apple.coreml.scripts.export --model_name="$MODEL_NAME"
 python3 -m examples.xnnpack.aot_compiler --model_name="$MODEL_NAME" --delegate
 
+# No generic examples/ entry point takes a model name and lowers it to MLX, so do
+# it here. Assert MLX claimed the graph before writing: the partitioner only warns
+# when it delegates nothing. This needs the MLX backend built, which the installer
+# does only on Apple Silicon with the Metal compiler present.
+python3 - "$MODEL_NAME" <<'PY'
+import sys
+
+import torch
+from executorch.backends.mlx import MLXPartitioner
+from executorch.examples.models import MODEL_NAME_TO_MODEL
+from executorch.examples.models.model_factory import EagerModelFactory
+from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower
+
+model_name = sys.argv[1]
+model, example_inputs, _, _ = EagerModelFactory.create_model(
+ *MODEL_NAME_TO_MODEL[model_name]
+)
+program = to_edge_transform_and_lower(
+ torch.export.export(model.eval(), example_inputs),
+ partitioner=[MLXPartitioner()],
+ compile_config=EdgeCompileConfig(_skip_dim_order=True),
+).to_executorch()
+
+segments = sum(
+ delegate.id == "MLXBackend"
+ for plan in program.executorch_program.execution_plan
+ for delegate in plan.delegates
+)
+if segments == 0:
+ raise SystemExit(f"error: {model_name} produced no MLX delegate")
+
+path = f"{model_name}_mlx.pte"
+with open(path, "wb") as file:
+ program.write_to_file(file)
+print(f"{path}: {segments} MLX delegate segment(s)")
+PY
+
 mkdir -p "$APP_PATH/Resources/Models/MobileNet/"
 mv $MODEL_NAME*.pte "$APP_PATH/Resources/Models/MobileNet/"