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
coremltoolssupports up to python 3.13. Do not run hatch with a newer version. Can be controlled using fxexport HATCH_PYTHON=python3.13- Run tests using
hatch run test:pytest tests
1 maintainer
Metadata
1 maintainer
Release files for stablehlo-coreml-experimental 0.0.18
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| stablehlo_coreml_experimental-0.0.18.tar.gz | 47.9 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| stablehlo_coreml_experimental-0.0.18-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 80.6 kB
Release files / stablehlo_coreml_experimental-0.0.18.tar.gz
| Download URL | stablehlo_coreml_experimental-0.0.18.tar.gz |
|---|---|
| Size | 47.9 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
53616eed4fde1017aa43b5ec1785863e60005588456532eca09b556451f63f03
|
|
BLAKE2b-256 checksum How to use checksums |
a6a2deca6b0e9b728b3aa06ff89d4ae09926945ef4ed5df79cd0838fcbdb4a6b
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
Hatch/1.16.5 cpython/3.12.3 HTTPX/0.28.1
|
Release files / stablehlo_coreml_experimental-0.0.18-py3-none-any.whl
| Download URL | stablehlo_coreml_experimental-0.0.18-py3-none-any.whl |
|---|---|
| Size | 32.7 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
cf4c527ff44d1ba905a73caf21d58185369e32ef0e318691d02b95df163738b5
|
|
BLAKE2b-256 checksum How to use checksums |
992d5bbe530a8604831b6f412ac69825c8a82f07e8b19e844fcf39cf92c228d0
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
Hatch/1.16.5 cpython/3.12.3 HTTPX/0.28.1
|
