Convert StableHLO models into Apple Core ML format

This repo is currently experimental!

Only a subset of the StableHLO operations have been implemented, and some of them may have restrictions.

Due to the current dot_general op implementation, it is only possible to target iOS >= 18.

Look in the tests directory, to see what has currently been tested.

The package is published to PyPi as stablehlo-coreml-experimental.

Converting a model

To convert a StableHLO module, do the following:

import coremltools as ct
from stablehlo_coreml.converter import convert
from stablehlo_coreml import DEFAULT_HLO_PIPELINE

mil_program = convert(hlo_module, minimum_deployment_target=ct.target.iOS18)
cml_model = ct.convert(
    mil_program,
    source="milinternal",
    minimum_deployment_target=ct.target.iOS18,
    pass_pipeline=DEFAULT_HLO_PIPELINE,
)

For a Jax project, the hlo_module can be obtained the following way:

import jax
from jax._src.lib.mlir import ir
from jax._src.interpreters import mlir as jax_mlir
from jax.export import export

import jax.numpy as jnp

def jax_function(a, b):
    return jnp.einsum("ij,jk -> ik", a, b)

context = jax_mlir.make_ir_context()
input_shapes = (jnp.zeros((2, 4)), jnp.zeros((4, 3)))
jax_exported = export(jax.jit(jax_function))(*input_shapes)
hlo_module = ir.Module.parse(jax_exported.mlir_module(), context=context)

For the Jax example to work, you will additionally need to install absl-py and flatbuffers as dependencies.

For additional examples see the tests directory.

Notes

• coremltools supports up to python 3.11. Do not run hatch with a newer version. Can be controlled using fx export HATCH_PYTHON=python3.11
• Run tests using hatch run test:pytest tests