Function is primarily designed to compile AI inference functions to run on-device. We will walk through the general workflow required to compile these functions.

​
Defining an AI Function

Let’s begin with a function that classifies an image, returning the label along with a confidence score. To do so, we will use the MobileNet v2 model from torchvision:

ai.py
from PIL import Image
from torch import argmax, inference_mode, softmax, randn
from torchvision.models import mobilenet_v2, MobileNet_V2_Weights
from torchvision.transforms import functional as F

weights = MobileNet_V2_Weights.DEFAULT
model = mobilenet_v2(weights=weights)
model.eval()

@inference_mode()
def predict (image: Image.Image) -> tuple[str, float]:
    """Classify an image."""
    # Preprocess
    image = image.convert("RGB")
    image = F.resize(image, 224)
    image = F.center_crop(image, 224)
    image_tensor = F.to_tensor(image)
    normalized_tensor = F.normalize(
        image_tensor,
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
    # Run model
    logits = model(normalized_tensor[None])
    scores = softmax(logits, dim=1)
    idx = argmax(scores, dim=1)
    score = scores[0,idx].item()
    label = weights.meta["categories"][idx]
    # Return
    return label, score

The code above has nothing that is specific to Function. It is plain PyTorch code.

​
Compiling the AI Function

There are a few steps needed to prepare an AI function for compilation:

In this section, required changes to the above code are highlighted.

​
Decorating the Function

First, apply the @compile decorator to the function to prepare it for compilation:

ai.py
from fxn import compile
from PIL import Image
from torch import argmax, inference_mode, softmax, randn
from torchvision.models import mobilenet_v2, MobileNet_V2_Weights
from torchvision.transforms import functional as F

weights = MobileNet_V2_Weights.DEFAULT
model = mobilenet_v2(weights=weights)
model.eval()

@compile(
    tag="@yusuf/classify-image",
    description="Classify an image with AI."
)
@inference_mode()
def predict (image: Image.Image) -> tuple[str, float]:
    """Classify an image."""
    # Preprocess
    image = image.convert("RGB")
    image = F.resize(image, 224)
    image = F.center_crop(image, 224)
    image_tensor = F.to_tensor(image)
    normalized_tensor = F.normalize(
        image_tensor,
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
    # Run model
    logits = model(normalized_tensor[None])
    scores = softmax(logits, dim=1)
    idx = argmax(scores, dim=1)
    score = scores[0,idx].item()
    label = weights.meta["categories"][idx]
    # Return
    return label, score

​
Defining the Compiler Sandbox

Depending on how you run AI inference, you will likely have to install libraries (e.g. PyTorch) and/or upload model weights. To do so, create a Sandbox:

ai.py
from fxn import compile, Sandbox
from PIL import Image
from torch import argmax, inference_mode, softmax, randn
from torchvision.models import mobilenet_v2, MobileNet_V2_Weights
from torchvision.transforms import functional as F

weights = MobileNet_V2_Weights.DEFAULT
model = mobilenet_v2(weights=weights)
model.eval()

@compile(
    tag="@yusuf/classify-image",
    description="Classify an image with AI.",
    sandbox=Sandbox().pip_install("torch", "torchvision")
)
@inference_mode()
def predict (image: Image.Image) -> tuple[str, float]:
    """Classify an image."""
    # Preprocess
    image = image.convert("RGB")
    image = F.resize(image, 224)
    image = F.center_crop(image, 224)
    image_tensor = F.to_tensor(image)
    normalized_tensor = F.normalize(
        image_tensor,
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
    # Run model
    logits = model(normalized_tensor[None])
    scores = softmax(logits, dim=1)
    idx = argmax(scores, dim=1)
    score = scores[0,idx].item()
    label = weights.meta["categories"][idx]
    # Return
    return label, score

​
Specifying Inference Backends

Let’s use the ONNXRuntime inference backend to run the AI model:

ai.py
from fxn import compile, Sandbox
from fxn.beta import ONNXInferenceMetadata
from PIL import Image
from torch import argmax, inference_mode, softmax, randn
from torchvision.models import mobilenet_v2, MobileNet_V2_Weights
from torchvision.transforms import functional as F

weights = MobileNet_V2_Weights.DEFAULT
model = mobilenet_v2(weights=weights)
model.eval()

@compile(
    tag="@yusuf/classify-image",
    description="Classify an image with AI.",
    sandbox=Sandbox().pip_install("torch", "torchvision"),
    metadata=[
      ONNXInferenceMetadata(
        model=model,
        model_args=(randn(1,3,224,224),)
      )
    ]
)
@inference_mode()
def predict (image: Image.Image) -> tuple[str, float]:
    """Classify an image."""
    # Preprocess
    image = image.convert("RGB")
    image = F.resize(image, 224)
    image = F.center_crop(image, 224)
    image_tensor = F.to_tensor(image)
    normalized_tensor = F.normalize(
        image_tensor,
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
    # Run model
    logits = model(normalized_tensor[None])
    scores = softmax(logits, dim=1)
    idx = argmax(scores, dim=1)
    score = scores[0,idx].item()
    label = weights.meta["categories"][idx]
    # Return
    return label, score

​
Compiling the Function

Now, compile the function using the Function CLI:

# Compile the AI function
$ fxn compile --overwrite ai.py

​
Supported Inference Backends

Function supports a fixed set of backends for running AI inference. You must opt in to using an inference backend for a specific model by providing inference metadata. The provided metadata will allow the Function compiler to lower the inference operation to native code. Below are supported inference metadata types:

You can use multiple inference backends to run a single model by providing multiple metadata instances that refer to the same model.

​
CoreML

Use the CoreMLInferenceMetadata metadata type to compile a PyTorch model to CoreML:

ai.py
from fxn.beta import CoreMLInferenceMetadata

# Create CoreML inference metadata
metadata = CoreMLInferenceMetadata(
    model=pytorch_model,
    model_args=(example_input,)
)

The CoreML inference backend is available only on iOS, macOS, and visionOS.

​
ONNX

Use the ONNXInferenceMetadata metadata type to compile a Python model to ONNX:

ai.py
from fxn.beta import ONNXInferenceMetadata

# Create ONNX inference metadata
metadata = ONNXInferenceMetadata(
    model=pytorch_model,
    model_args=(example_input,)
)

​
ONNX Runtime

Use the ONNXRuntimeInferenceSessionMetadata metadata type to compile an existing ONNX model and run it with ONNXRuntime:

ai.py
from fxn.beta import ONNXRuntimeInferenceSessionMetadata

# Create ONNXRuntime inference metadata
metadata = ONNXRuntimeInferenceSessionMetadata(
    session=ort_session,
    model_path="/path/to/model.onnx"
)

The model must exist at the provided model_path in the compiler sandbox.

​
Llama.cpp

Coming soon 🤫.