import json
import pathlib
import sys
import tempfile
import coremltools as ct
import numpy as np
from coremltools.converters.mil import Builder as mb
from coremltools.converters.mil.mil import types
COREMLTOOLS = "9.0"
def single_function_package(directory, name, size, factor, output):
@mb.program(
input_specs=[mb.TensorSpec(shape=(size,), dtype=types.fp32)],
opset_version=ct.target.macOS15,
)
def program(x):
return mb.mul(x=x, y=np.float32(factor), name=output)
model = ct.convert(
program,
convert_to="mlprogram",
minimum_deployment_target=ct.target.macOS15,
compute_precision=ct.precision.FLOAT32,
)
path = pathlib.Path(directory) / f"{name}.mlpackage"
model.save(str(path))
return path
def drop_empty_weights(package):
weights = package / "Data" / "com.apple.CoreML" / "weights"
if any(weights.iterdir()):
sys.exit(f"{weights} is not empty; the programs were meant to carry no weights")
weights.rmdir()
manifest_path = package / "Manifest.json"
manifest = json.loads(manifest_path.read_text())
items = manifest["itemInfoEntries"]
dropped = [key for key, item in items.items() if item["path"] == "com.apple.CoreML/weights"]
if len(dropped) != 1:
sys.exit(f"expected one `weights` item in {manifest_path}, found {len(dropped)}")
del items[dropped[0]]
manifest_path.write_text(json.dumps(manifest, indent=4))
def main():
if len(sys.argv) != 2:
sys.exit("usage: python -I two_functions.py <out-dir>")
if ct.__version__ != COREMLTOOLS:
sys.exit(f"coremltools {ct.__version__} found; this fixture is built with {COREMLTOOLS}")
out = pathlib.Path(sys.argv[1]) / "two_functions.mlpackage"
if out.exists():
sys.exit(f"{out} exists; remove it first")
with tempfile.TemporaryDirectory() as scratch:
double = single_function_package(scratch, "double", 2, 2.0, "y")
negate = single_function_package(scratch, "negate", 3, -1.0, "z")
descriptor = ct.utils.MultiFunctionDescriptor()
descriptor.add_function(str(double), src_function_name="main", target_function_name="double")
descriptor.add_function(str(negate), src_function_name="main", target_function_name="negate")
descriptor.default_function_name = "negate"
ct.utils.save_multifunction(descriptor, str(out))
drop_empty_weights(out)
print(out)
if __name__ == "__main__":
main()