export_functions(functions: dict[str, tuple[Module, object]], path: str, default: str | None = None, precision: Precision = Precision.FLOAT32, weights: WeightPrecision | Palettize | Sparsify = WeightPrecision.FLOAT, metadata: Metadata | None = None, compute_units: ComputeUnits = ComputeUnits.ALL)Write several entry points into one package, sharing its weights.
A decoder wants two: one that reads a whole prompt and one that reads a single token. They are the same network, and shipping them as two packages ships the weights twice. Here a parameter is written once and every function that uses it points at the same bytes.
Needs iOS 18 / macOS 15, which is what several entry points costs.
Parameters
functionsdict[str, tuple[nn.Module, object]]Name to the model and the example input to trace it with. The
models may be the same object; that is the case worth having.
pathstrDestination
.mlpackage. Replaced if it exists.default(str or None, optional, keyword - only)= NoneEntry point a caller gets without asking. The first, if unnamed.
Body precision, for every function.
Weight storage, for every function.
What the package says about itself. One package, one set.
Which processors Core ML may schedule on.
Returns
dict[str, CoreMLModel]One handle per function, each pinned to its own entry point.
Raises
ValueErrorNo functions, or
default names one that is not there.Examples
A decoder wants two entry points: one that reads a whole prompt and
one that reads a single token. They are the same network, so the
weights are written once and both point at the same bytes.
>>> import shutil, tempfile
>>> import lucid, lucid.nn as nn, lucid.coreml as cml
>>> class Decoder(nn.Module):
... def __init__(self):
... super().__init__()
... self.embed = nn.Embedding(100, 16)
... self.head = nn.Linear(16, 100)
... def forward(self, ids):
... return self.head(self.embed(ids))
>>> model, room = Decoder().eval(), tempfile.mkdtemp()
>>> handles = cml.export_functions(
... {
... "prompt": (model, lucid.zeros(1, 8).to(lucid.int64)),
... "step": (model, lucid.zeros(1, 1).to(lucid.int64)),
... },
... f"{room}/decoder.mlpackage",
... default="step",
... )
>>> prompt_ids = lucid.randint(0, 100, (1, 8))
>>> handles["prompt"].predict(prompt_ids).shape
(1, 8, 100)
>>> handles["step"].predict(prompt_ids[:, -1:]).shape
(1, 1, 100)
>>> for handle in handles.values():
... handle.close()
>>> shutil.rmtree(room)