From 283b057dd69ae4f3ef67beb1f03409330b565662 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 24 Mar 2026 16:39:06 +0000 Subject: [PATCH 1/5] Initial plan From 9e8ebe1c0ffd72d47ee237167b1346a6174dbc0d Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 24 Mar 2026 16:49:01 +0000 Subject: [PATCH 2/5] Remove onnxruntime-easy dependency, inline its code into ort_inference.py Co-authored-by: xadupre <22452781+xadupre@users.noreply.github.com> Agent-Logs-Url: https://github.com/onnxruntime/mobius/sessions/2a25ff2f-35b5-4acf-b7e5-4f45dbca8aee --- pyproject.toml | 3 +- src/mobius/_testing/ort_inference.py | 102 ++++++++++++++++++++++++--- 2 files changed, 96 insertions(+), 9 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2d6b1813..1c10208d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,7 +28,8 @@ mobius = "mobius.__main__:main" transformers = ["transformers", "safetensors"] gguf = ["gguf>=0.10.0"] testing = [ - "onnxruntime-easy", + "onnxruntime", + "ml_dtypes", "torch", "transformers>=5.0", "safetensors", diff --git a/src/mobius/_testing/ort_inference.py b/src/mobius/_testing/ort_inference.py index 5d4665bd..1f81a2e0 100644 --- a/src/mobius/_testing/ort_inference.py +++ b/src/mobius/_testing/ort_inference.py @@ -3,24 +3,105 @@ """ONNX Runtime inference session wrapper for ir.Model objects. -Uses ``onnxruntime-easy`` which handles bfloat16 and other non-standard -dtypes transparently. +Handles bfloat16 and other non-standard dtypes transparently via ml_dtypes. """ from __future__ import annotations +import importlib.util import tempfile from pathlib import Path import numpy as np import onnx_ir as ir -import onnxruntime_easy as ort_easy +import onnxruntime as ort +import onnxruntime.capi._pybind_state as _ort_c from mobius._model_package import ModelPackage +_HAS_ML_DTYPES = importlib.util.find_spec("ml_dtypes") is not None +if _HAS_ML_DTYPES: + # ml_dtypes provides bfloat16 and float8 variants not in standard numpy. + # When absent, OrtValues are created via the plain numpy path and these + # dtypes are unsupported. + import ml_dtypes + +# ONNX element type constants for dtypes not representable as plain numpy dtypes +_BFLOAT16_TYPE = 16 +_FLOAT8E4M3FN_TYPE = 17 +_FLOAT8E4M3FNUZ_TYPE = 18 +_FLOAT8E5M2_TYPE = 19 +_FLOAT8E5M2FNUZ_TYPE = 20 +_UINT4_TYPE = 21 +_INT4_TYPE = 22 +_FLOAT4E2M1_TYPE = 23 + + +def _ml_dtype_to_onnx_type(dtype: np.dtype) -> int | None: + """Return the ONNX element type integer for an ml_dtypes dtype, or None. + + Only call this function when ``_HAS_ML_DTYPES`` is True. + """ + if not _HAS_ML_DTYPES: + return None + if dtype == ml_dtypes.bfloat16: + return _BFLOAT16_TYPE + if dtype == ml_dtypes.float8_e4m3fn: + return _FLOAT8E4M3FN_TYPE + if dtype == ml_dtypes.float8_e4m3fnuz: + return _FLOAT8E4M3FNUZ_TYPE + if dtype == ml_dtypes.float8_e5m2: + return _FLOAT8E5M2_TYPE + if dtype == ml_dtypes.float8_e5m2fnuz: + return _FLOAT8E5M2FNUZ_TYPE + if dtype == ml_dtypes.uint4: + return _UINT4_TYPE + if dtype == ml_dtypes.int4: + return _INT4_TYPE + if dtype == ml_dtypes.float4_e2m1fn: + return _FLOAT4E2M1_TYPE + return None + + +def _to_ort_value(value: np.ndarray, device: str = "cpu") -> ort.OrtValue: + """Convert a numpy array to an OrtValue, handling special ml_dtypes dtypes.""" + # Use DLPack when available (e.g. torch tensors or non-zero-size arrays) + if hasattr(value, "__dlpack__"): + is_zero_size = hasattr(value, "size") and value.size == 0 + if not is_zero_size: + return ort.OrtValue( + _ort_c.OrtValue.from_dlpack(value.__dlpack__(), False), value + ) + if _HAS_ML_DTYPES and isinstance(value, np.ndarray): + onnx_type = _ml_dtype_to_onnx_type(value.dtype) + if onnx_type is not None: + return ort.OrtValue.ortvalue_from_numpy_with_onnx_type( + value, onnx_element_type=onnx_type + ) + return ort.OrtValue.ortvalue_from_numpy(np.asarray(value), device) + + +def _create_session(model_path: str, device: str = "cpu") -> ort.InferenceSession: + """Create an ORT InferenceSession with sensible defaults. + + Args: + model_path: Path to the ONNX model file. + device: Execution device, ``"cpu"`` or ``"cuda"``. + """ + if device == "cpu": + providers = ("CPUExecutionProvider",) + elif device == "cuda": + providers = ("CUDAExecutionProvider", "CPUExecutionProvider") + else: + raise ValueError(f"Unsupported device: {device!r}. Expected 'cpu' or 'cuda'.") + opts = ort.SessionOptions() + opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL + opts.log_severity_level = 2 # warning + return ort.InferenceSession(model_path, sess_options=opts, providers=providers) + class OnnxModelSession: - """Wraps an ``onnxruntime_easy.EasySession`` for an ``ir.Model``. + """Wraps an ``ort.InferenceSession`` for an ``ir.Model``. Serializes the model to a temporary file and creates an ORT session. Provides a simple ``run()`` interface that accepts and returns numpy arrays. @@ -37,7 +118,7 @@ class OnnxModelSession: def __init__( self, model: ir.Model | ModelPackage, - **load_kwargs, + device: str = "cpu", ): if isinstance(model, ModelPackage): if len(model) != 1: @@ -46,11 +127,12 @@ def __init__( f"single ir.Model or index into the package." ) model = next(iter(model.values())) + self._device = device self._tmpdir = tempfile.TemporaryDirectory() self._model_path = str(Path(self._tmpdir.name) / "model.onnx") ir.save(model, self._model_path, external_data="model.onnx.data") - self._session = ort_easy.load(self._model_path, **load_kwargs) + self._session = _create_session(self._model_path, device=device) self._input_names = [inp.name for inp in self._session.get_inputs()] self._output_names = [out.name for out in self._session.get_outputs()] @@ -81,8 +163,12 @@ def run(self, feeds: dict[str, np.ndarray]) -> dict[str, np.ndarray]: # because np.ascontiguousarray promotes them to 1-d. if v.ndim > 0: v = np.ascontiguousarray(v) - ort_feeds[k] = ort_easy.ort_value(v) - raw_outputs = self._session(**ort_feeds) + ort_feeds[k] = _to_ort_value(v, device=self._device) + run_opts = ort.RunOptions() + run_opts.log_severity_level = 2 # warning + raw_outputs = self._session.run_with_ort_values( + None, ort_feeds, run_options=run_opts + ) return dict(zip(self._output_names, (o.numpy() for o in raw_outputs))) def close(self) -> None: From f2219898eb525ff9865f8f216e8cbbd5fbf81df0 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 24 Mar 2026 16:56:14 +0000 Subject: [PATCH 3/5] Use onnx.TensorProto constants instead of hardcoded integers Co-authored-by: xadupre <22452781+xadupre@users.noreply.github.com> Agent-Logs-Url: https://github.com/onnxruntime/mobius/sessions/60380e38-d669-4ed2-b8cd-d74d9ee42c9c --- src/mobius/_testing/ort_inference.py | 28 ++++++++++------------------ 1 file changed, 10 insertions(+), 18 deletions(-) diff --git a/src/mobius/_testing/ort_inference.py b/src/mobius/_testing/ort_inference.py index 1f81a2e0..c7208630 100644 --- a/src/mobius/_testing/ort_inference.py +++ b/src/mobius/_testing/ort_inference.py @@ -13,6 +13,7 @@ from pathlib import Path import numpy as np +import onnx import onnx_ir as ir import onnxruntime as ort import onnxruntime.capi._pybind_state as _ort_c @@ -26,16 +27,6 @@ # dtypes are unsupported. import ml_dtypes -# ONNX element type constants for dtypes not representable as plain numpy dtypes -_BFLOAT16_TYPE = 16 -_FLOAT8E4M3FN_TYPE = 17 -_FLOAT8E4M3FNUZ_TYPE = 18 -_FLOAT8E5M2_TYPE = 19 -_FLOAT8E5M2FNUZ_TYPE = 20 -_UINT4_TYPE = 21 -_INT4_TYPE = 22 -_FLOAT4E2M1_TYPE = 23 - def _ml_dtype_to_onnx_type(dtype: np.dtype) -> int | None: """Return the ONNX element type integer for an ml_dtypes dtype, or None. @@ -44,22 +35,23 @@ def _ml_dtype_to_onnx_type(dtype: np.dtype) -> int | None: """ if not _HAS_ML_DTYPES: return None + tp = onnx.TensorProto if dtype == ml_dtypes.bfloat16: - return _BFLOAT16_TYPE + return tp.BFLOAT16 if dtype == ml_dtypes.float8_e4m3fn: - return _FLOAT8E4M3FN_TYPE + return tp.FLOAT8E4M3FN if dtype == ml_dtypes.float8_e4m3fnuz: - return _FLOAT8E4M3FNUZ_TYPE + return tp.FLOAT8E4M3FNUZ if dtype == ml_dtypes.float8_e5m2: - return _FLOAT8E5M2_TYPE + return tp.FLOAT8E5M2 if dtype == ml_dtypes.float8_e5m2fnuz: - return _FLOAT8E5M2FNUZ_TYPE + return tp.FLOAT8E5M2FNUZ if dtype == ml_dtypes.uint4: - return _UINT4_TYPE + return tp.UINT4 if dtype == ml_dtypes.int4: - return _INT4_TYPE + return tp.INT4 if dtype == ml_dtypes.float4_e2m1fn: - return _FLOAT4E2M1_TYPE + return tp.FLOAT4E2M1 return None From 439ecbbd92a04076fa77ca1427bf8debe0f9a129 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 24 Mar 2026 17:36:34 +0000 Subject: [PATCH 4/5] Simplify _to_ort_value: use ir.DataType.from_numpy, drop ml_dtypes conditional Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com> Agent-Logs-Url: https://github.com/onnxruntime/mobius/sessions/feb5cf18-6cf6-432e-8ef5-ac0224beb18a --- src/mobius/_testing/ort_inference.py | 47 +++++----------------------- 1 file changed, 7 insertions(+), 40 deletions(-) diff --git a/src/mobius/_testing/ort_inference.py b/src/mobius/_testing/ort_inference.py index c7208630..1b8f3d33 100644 --- a/src/mobius/_testing/ort_inference.py +++ b/src/mobius/_testing/ort_inference.py @@ -8,52 +8,16 @@ from __future__ import annotations -import importlib.util import tempfile from pathlib import Path import numpy as np -import onnx import onnx_ir as ir import onnxruntime as ort import onnxruntime.capi._pybind_state as _ort_c from mobius._model_package import ModelPackage -_HAS_ML_DTYPES = importlib.util.find_spec("ml_dtypes") is not None -if _HAS_ML_DTYPES: - # ml_dtypes provides bfloat16 and float8 variants not in standard numpy. - # When absent, OrtValues are created via the plain numpy path and these - # dtypes are unsupported. - import ml_dtypes - - -def _ml_dtype_to_onnx_type(dtype: np.dtype) -> int | None: - """Return the ONNX element type integer for an ml_dtypes dtype, or None. - - Only call this function when ``_HAS_ML_DTYPES`` is True. - """ - if not _HAS_ML_DTYPES: - return None - tp = onnx.TensorProto - if dtype == ml_dtypes.bfloat16: - return tp.BFLOAT16 - if dtype == ml_dtypes.float8_e4m3fn: - return tp.FLOAT8E4M3FN - if dtype == ml_dtypes.float8_e4m3fnuz: - return tp.FLOAT8E4M3FNUZ - if dtype == ml_dtypes.float8_e5m2: - return tp.FLOAT8E5M2 - if dtype == ml_dtypes.float8_e5m2fnuz: - return tp.FLOAT8E5M2FNUZ - if dtype == ml_dtypes.uint4: - return tp.UINT4 - if dtype == ml_dtypes.int4: - return tp.INT4 - if dtype == ml_dtypes.float4_e2m1fn: - return tp.FLOAT4E2M1 - return None - def _to_ort_value(value: np.ndarray, device: str = "cpu") -> ort.OrtValue: """Convert a numpy array to an OrtValue, handling special ml_dtypes dtypes.""" @@ -64,11 +28,14 @@ def _to_ort_value(value: np.ndarray, device: str = "cpu") -> ort.OrtValue: return ort.OrtValue( _ort_c.OrtValue.from_dlpack(value.__dlpack__(), False), value ) - if _HAS_ML_DTYPES and isinstance(value, np.ndarray): - onnx_type = _ml_dtype_to_onnx_type(value.dtype) - if onnx_type is not None: + if isinstance(value, np.ndarray): + try: + onnx_type = ir.DataType.from_numpy(value.dtype) + except TypeError: + pass + else: return ort.OrtValue.ortvalue_from_numpy_with_onnx_type( - value, onnx_element_type=onnx_type + value, onnx_element_type=onnx_type.value ) return ort.OrtValue.ortvalue_from_numpy(np.asarray(value), device) From 055d5a70fe373ed245d81bd610e43c3e9da22219 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Xavier=20Dupr=C3=A9?= Date: Wed, 25 Mar 2026 11:00:07 +0100 Subject: [PATCH 5/5] lint --- src/mobius/_testing/ort_inference.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/src/mobius/_testing/ort_inference.py b/src/mobius/_testing/ort_inference.py index 1b8f3d33..ea3e5c93 100644 --- a/src/mobius/_testing/ort_inference.py +++ b/src/mobius/_testing/ort_inference.py @@ -25,9 +25,7 @@ def _to_ort_value(value: np.ndarray, device: str = "cpu") -> ort.OrtValue: if hasattr(value, "__dlpack__"): is_zero_size = hasattr(value, "size") and value.size == 0 if not is_zero_size: - return ort.OrtValue( - _ort_c.OrtValue.from_dlpack(value.__dlpack__(), False), value - ) + return ort.OrtValue(_ort_c.OrtValue.from_dlpack(value.__dlpack__(), False), value) if isinstance(value, np.ndarray): try: onnx_type = ir.DataType.from_numpy(value.dtype) @@ -125,9 +123,7 @@ def run(self, feeds: dict[str, np.ndarray]) -> dict[str, np.ndarray]: ort_feeds[k] = _to_ort_value(v, device=self._device) run_opts = ort.RunOptions() run_opts.log_severity_level = 2 # warning - raw_outputs = self._session.run_with_ort_values( - None, ort_feeds, run_options=run_opts - ) + raw_outputs = self._session.run_with_ort_values(None, ort_feeds, run_options=run_opts) return dict(zip(self._output_names, (o.numpy() for o in raw_outputs))) def close(self) -> None: