# 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:
```python
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:
```python
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.12. Do not run hatch with a newer version.
Can be controlled using fx `export HATCH_PYTHON=python3.12`
* Run tests using `hatch run test:pytest tests`
Raw data
{
"_id": null,
"home_page": null,
"name": "stablehlo-coreml-experimental",
"maintainer": null,
"docs_url": null,
"requires_python": ">=3.9",
"maintainer_email": null,
"keywords": "converter, coreml, coremltools, hlo, machinelearning, ml, neural, stablehlo, xla",
"author": null,
"author_email": "Kasper Nielsen <kasper0406@gmail.com>",
"download_url": "https://files.pythonhosted.org/packages/99/b5/359d71f2ea6e32859fbc26ca84e2815a86db0c7a59857307ca620f814973/stablehlo_coreml_experimental-0.0.11.tar.gz",
"platform": null,
"description": "# Convert StableHLO models into Apple Core ML format\n\n**This repo is currently experimental!**\n\nOnly a subset of the StableHLO operations have been implemented, and some of them may have restrictions.\n\nDue to the current _dot_general_ op implementation, it is only possible to target iOS >= 18.\n\nLook in the `tests` directory, to see what has currently been tested.\n\nThe package is published to PyPi as `stablehlo-coreml-experimental`.\n\n## Converting a model\n\nTo convert a StableHLO module, do the following:\n\n```python\nimport coremltools as ct\nfrom stablehlo_coreml.converter import convert\nfrom stablehlo_coreml import DEFAULT_HLO_PIPELINE\n\nmil_program = convert(hlo_module, minimum_deployment_target=ct.target.iOS18)\ncml_model = ct.convert(\n mil_program,\n source=\"milinternal\",\n minimum_deployment_target=ct.target.iOS18,\n pass_pipeline=DEFAULT_HLO_PIPELINE,\n)\n```\n\nFor a Jax project, the `hlo_module` can be obtained the following way:\n\n```python\nimport jax\nfrom jax._src.lib.mlir import ir\nfrom jax._src.interpreters import mlir as jax_mlir\nfrom jax.export import export\n\nimport jax.numpy as jnp\n\ndef jax_function(a, b):\n return jnp.einsum(\"ij,jk -> ik\", a, b)\n\ncontext = jax_mlir.make_ir_context()\ninput_shapes = (jnp.zeros((2, 4)), jnp.zeros((4, 3)))\njax_exported = export(jax.jit(jax_function))(*input_shapes)\nhlo_module = ir.Module.parse(jax_exported.mlir_module(), context=context)\n```\n\nFor the Jax example to work, you will additionally need to install `absl-py` and `flatbuffers` as dependencies.\n\nFor additional examples see the `tests` directory.\n\n## Notes\n* `coremltools` supports up to python 3.12. Do not run hatch with a newer version.\n Can be controlled using fx `export HATCH_PYTHON=python3.12`\n* Run tests using `hatch run test:pytest tests`\n",
"bugtrack_url": null,
"license": null,
"summary": "Convert StableHLO models into Apple Core ML format",
"version": "0.0.11",
"project_urls": {
"Homepage": "https://github.com/kasper0406/stablehlo-coreml",
"Issues": "https://github.com/kasper0406/stablehlo-coreml/issues"
},
"split_keywords": [
"converter",
" coreml",
" coremltools",
" hlo",
" machinelearning",
" ml",
" neural",
" stablehlo",
" xla"
],
"urls": [
{
"comment_text": null,
"digests": {
"blake2b_256": "fadd5c58e64c96a270bee58275d2d2f045c326ce7271867558ab57e3d0c1a163",
"md5": "b61807283795f0b283826e4f1c03fcf6",
"sha256": "2dc8e6e7fd02701a86b25a2180b406bb13eb8d55e382b863f8a033f81274295a"
},
"downloads": -1,
"filename": "stablehlo_coreml_experimental-0.0.11-py3-none-any.whl",
"has_sig": false,
"md5_digest": "b61807283795f0b283826e4f1c03fcf6",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": ">=3.9",
"size": 21509,
"upload_time": "2024-11-20T18:25:19",
"upload_time_iso_8601": "2024-11-20T18:25:19.731867Z",
"url": "https://files.pythonhosted.org/packages/fa/dd/5c58e64c96a270bee58275d2d2f045c326ce7271867558ab57e3d0c1a163/stablehlo_coreml_experimental-0.0.11-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": null,
"digests": {
"blake2b_256": "99b5359d71f2ea6e32859fbc26ca84e2815a86db0c7a59857307ca620f814973",
"md5": "e782313ea8c70676cbdb50e481818a84",
"sha256": "a1a70cbc05202d8aadee25d5c62832254e4aea94e8dd997b54f04fd8762eb24f"
},
"downloads": -1,
"filename": "stablehlo_coreml_experimental-0.0.11.tar.gz",
"has_sig": false,
"md5_digest": "e782313ea8c70676cbdb50e481818a84",
"packagetype": "sdist",
"python_version": "source",
"requires_python": ">=3.9",
"size": 33861,
"upload_time": "2024-11-20T18:25:18",
"upload_time_iso_8601": "2024-11-20T18:25:18.866148Z",
"url": "https://files.pythonhosted.org/packages/99/b5/359d71f2ea6e32859fbc26ca84e2815a86db0c7a59857307ca620f814973/stablehlo_coreml_experimental-0.0.11.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2024-11-20 18:25:18",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "kasper0406",
"github_project": "stablehlo-coreml",
"travis_ci": false,
"coveralls": false,
"github_actions": true,
"lcname": "stablehlo-coreml-experimental"
}
JSON
