onnx-diagnostic 0.8.10__py3-none-any.whl → 0.8.11__py3-none-any.whl

onnx_diagnostic/helpers/onnx_helper.py

@@ -1,6 +1,7 @@
 import functools
 import json
 import os
+import re
 import sys
 import warnings
 from typing import (
@@ -32,11 +33,10 @@ from onnx import (
     ValueInfoProto,
     load as onnx_load,
 )
+from ..typing import InferenceSessionLike, TensorLike
 
-TensorLike = Union[np.ndarray, "torch.Tensor"]  # noqa: F821
 
-
-def _make_stat(init: TensorProto) -> Dict[str, float]:
+def _make_stat(init: TensorProto) -> Dict[str, Any]:
     """
     Produces statistics.
 
@@ -160,11 +160,11 @@ def _validate_graph(
     verbose: int = 0,
     watch: Optional[Set[str]] = None,
     path: Optional[Sequence[str]] = None,
-):
-    found = []
+) -> List[Union[NodeProto, TensorProto, ValueInfoProto]]:
+    found: List[Union[NodeProto, TensorProto, ValueInfoProto]] = []
     path = path or ["root"]
-    set_init = set(i.name for i in g.initializer)
-    set_input = set(i.name for i in g.input)
+    set_init = {i.name for i in g.initializer}
+    set_input = {i.name for i in g.input}
     existing |= set_init | set_input
     if watch and set_init & watch:
         if verbose:
@@ -215,18 +215,15 @@ def _validate_graph(
                     f"in {'/'.join(path)}/{node.op_type}[{node.name}]"
                 )
             found.append(node)
-    out = set(o.name for o in g.output)
+    out = {o.name for o in g.output}
     ins = out & existing
-    if ins != out:
-        raise AssertionError(
-            f"One output is missing, out={node.input}, existing={ins}, path={path}"
-        )
+    assert ins == out, f"One output is missing, out={node.input}, existing={ins}, path={path}"
     return found
 
 
 def _validate_function(g: FunctionProto, verbose: int = 0, watch: Optional[Set[str]] = None):
-    existing = set(g.input)
-    found = []
+    existing: Set[str] = set(g.input)
+    found: List[Union[NodeProto, TensorProto, ValueInfoProto]] = []
     for node in g.node:
         ins = set(node.input) & existing
         if ins != set(node.input):
@@ -240,7 +237,7 @@ def _validate_function(g: FunctionProto, verbose: int = 0, watch: Optional[Set[s
         for att in node.attribute:
             if att.type == AttributeProto.GRAPH:
                 found.extend(
-                    _validate_graph(g, existing.copy(), path=[g.name], verbose=verbose)
+                    _validate_graph(att.g, existing.copy(), path=[g.name], verbose=verbose)
                 )
         existing |= set(node.output)
         if watch and set(node.output) & watch:
@@ -285,7 +282,7 @@ def check_model_ort(
     onx: ModelProto,
     providers: Optional[Union[str, List[Any]]] = None,
     dump_file: Optional[str] = None,
-) -> "onnxruntime.InferenceSession":  # noqa: F821
+) -> InferenceSessionLike:
     """
     Loads a model with onnxruntime.
 
@@ -308,10 +305,9 @@ def check_model_ort(
 
     if isinstance(onx, str):
         try:
+            # pyrefly: ignore[bad-return]
             return InferenceSession(onx, providers=providers)
         except Exception as e:
-            import onnx
-
             if dump_file:
                 onnx.save(onx, dump_file)
 
@@ -319,8 +315,8 @@ def check_model_ort(
                 f"onnxruntime cannot load the model "
                 f"due to {e}\n{pretty_onnx(onnx.load(onx))}"
             )
-        return
     try:
+        # pyrefly: ignore[bad-return]
         return InferenceSession(onx.SerializeToString(), providers=providers)
     except Exception as e:
         if dump_file:
@@ -358,7 +354,17 @@ def onnx_dtype_name(itype: int, exc: bool = True) -> str:
 
 
 def pretty_onnx(
-    onx: Union[FunctionProto, GraphProto, ModelProto, ValueInfoProto, str],
+    onx: Union[
+        AttributeProto,
+        FunctionProto,
+        GraphProto,
+        ModelProto,
+        NodeProto,
+        onnx.SparseTensorProto,
+        TensorProto,
+        ValueInfoProto,
+        str,
+    ],
     with_attributes: bool = False,
     highlight: Optional[Set[str]] = None,
     shape_inference: bool = False,
@@ -377,6 +383,9 @@ def pretty_onnx(
     assert onx is not None, "onx cannot be None"
 
     if shape_inference:
+        assert isinstance(
+            onx, ModelProto
+        ), f"shape inference only works for ModelProto, not {type(onx)}"
         onx = onnx.shape_inference.infer_shapes(onx)
 
     if isinstance(onx, ValueInfoProto):
@@ -447,6 +456,8 @@ def pretty_onnx(
         shape = "x".join(map(str, onx.dims))
         return f"TensorProto:{onx.data_type}:{shape}:{onx.name}"
 
+    assert not isinstance(onx, onnx.SparseTensorProto), "SparseTensorProto is not handled yet."
+
     try:
         from onnx_array_api.plotting.text_plot import onnx_simple_text_plot
 
@@ -538,12 +549,6 @@ def from_array_ml_dtypes(arr: TensorLike, name: Optional[str] = None) -> TensorP
     return tensor
 
 
-_STORAGE_TYPE = {
-    TensorProto.FLOAT16: np.int16,
-    TensorProto.BFLOAT16: np.int16,
-}
-
-
 def from_array_extended(tensor: TensorLike, name: Optional[str] = None) -> TensorProto:
     """
     Converts an array into a :class:`onnx.TensorProto`.
@@ -561,54 +566,9 @@ def from_array_extended(tensor: TensorLike, name: Optional[str] = None) -> Tenso
         ), f"Unable to convert type {type(tensor)} into TensorProto."
         return proto_from_tensor(tensor, name=name)
 
-    try:
-        from onnx.reference.ops.op_cast import (
-            bfloat16,
-            float8e4m3fn,
-            float8e4m3fnuz,
-            float8e5m2,
-            float8e5m2fnuz,
-        )
-    except ImportError:
-        bfloat16 = None
-
-    if bfloat16 is None:
-        return onh.from_array(tensor, name)
-
-    dt = tensor.dtype
-    if dt == float8e4m3fn and dt.descr[0][0] == "e4m3fn":
-        to = TensorProto.FLOAT8E4M3FN
-        dt_to = np.uint8
-    elif dt == float8e4m3fnuz and dt.descr[0][0] == "e4m3fnuz":
-        to = TensorProto.FLOAT8E4M3FNUZ
-        dt_to = np.uint8
-    elif dt == float8e5m2 and dt.descr[0][0] == "e5m2":
-        to = TensorProto.FLOAT8E5M2
-        dt_to = np.uint8
-    elif dt == float8e5m2fnuz and dt.descr[0][0] == "e5m2fnuz":
-        to = TensorProto.FLOAT8E5M2FNUZ
-        dt_to = np.uint8
-    elif dt == bfloat16 and dt.descr[0][0] == "bfloat16":
-        to = TensorProto.BFLOAT16
-        dt_to = np.uint16
-    else:
-        try:
-            import ml_dtypes
-        except ImportError:
-            ml_dtypes = None
-        if ml_dtypes is not None and (
-            tensor.dtype == ml_dtypes.bfloat16
-            or tensor.dtype == ml_dtypes.float8_e4m3fn
-            or tensor.dtype == ml_dtypes.float8_e4m3fnuz
-            or tensor.dtype == ml_dtypes.float8_e5m2
-            or tensor.dtype == ml_dtypes.float8_e5m2fnuz
-        ):
-            return from_array_ml_dtypes(tensor, name)
-        return onh.from_array(tensor, name)
-
-    t = onh.from_array(tensor.astype(dt_to), name)
-    t.data_type = to
-    return t
+    assert isinstance(tensor, np.ndarray)  # type checking
+    # pyrefly: ignore[bad-argument-type]
+    return onh.from_array(tensor, name)
 
 
 def to_array_extended(proto: TensorProto) -> TensorLike:
@@ -666,6 +626,7 @@ def onnx_dtype_to_np_dtype(itype: int) -> Any:
     )
 
 
+# pyrefly: ignore[unknown-name]
 def dtype_to_tensor_dtype(dt: Union[np.dtype, "torch.dtype"]) -> int:  # noqa: F821
     """
     Converts a torch dtype or numpy dtype into a onnx element type.
@@ -679,6 +640,7 @@ def dtype_to_tensor_dtype(dt: Union[np.dtype, "torch.dtype"]) -> int:  # noqa: F
         pass
     from .torch_helper import torch_dtype_to_onnx_dtype
 
+    # pyrefly: ignore[bad-argument-type]
     return torch_dtype_to_onnx_dtype(dt)
 
 
@@ -779,6 +741,7 @@ def tensor_dtype_to_np_dtype(tensor_dtype: int) -> np.dtype:
                 f"ml_dtypes can be used."
             ) from e
 
+        # pyrefly: ignore[bad-assignment]
         mapping: Dict[int, np.dtype] = {
             TensorProto.BFLOAT16: ml_dtypes.bfloat16,
             TensorProto.FLOAT8E4M3FN: ml_dtypes.float8_e4m3fn,
@@ -798,7 +761,7 @@ def iterator_initializer_constant(
     model: Union[FunctionProto, GraphProto, ModelProto],
     use_numpy: bool = True,
     prefix: str = "",
-) -> Iterator[Tuple[str, Union["torch.Tensor", np.ndarray]]]:  # noqa: F821
+) -> Iterator[Tuple[str, TensorLike]]:  # noqa: F821
     """
     Iterates on iniatialiers and constant in an onnx model.
 
@@ -814,9 +777,12 @@ def iterator_initializer_constant(
         if prefix:
             prefix += "."
         for init in graph.initializer:
-            yield f"{prefix}{init.name}", (
-                to_array_extended(init) if use_numpy else to_tensor(init)
-            )
+            s = f"{prefix}{init.name}"
+            if use_numpy:
+                yield s, to_array_extended(init)
+            else:
+                # pyrefly: ignore[unbound-name]
+                yield s, to_tensor(init)
         nodes = graph.node
         name = graph.name
         if isinstance(model, ModelProto):
@@ -831,13 +797,15 @@ def iterator_initializer_constant(
         if node.op_type == "Constant" and node.domain == "":
             from ..reference import ExtendedReferenceEvaluator as Inference
 
-            if not use_numpy:
-                import torch
             sess = Inference(node)
             value = sess.run(None, {})[0]
-            yield f"{prefix}{node.output[0]}", (
-                value if use_numpy else torch.from_numpy(value)
-            )
+
+            if not use_numpy:
+                import torch
+
+                yield f"{prefix}{node.output[0]}", (torch.from_numpy(value))
+            else:
+                yield f"{prefix}{node.output[0]}", (value)
 
         if node.op_type in {"Loop", "Body", "Scan"}:
             for att in node.attribute:
@@ -870,7 +838,9 @@ def tensor_statistics(tensor: Union[np.ndarray, TensorProto]) -> Dict[str, Union
     from .helper import size_type
 
     if isinstance(tensor, TensorProto):
+        # pyrefly: ignore[bad-assignment]
         tensor = to_array_extended(tensor)
+    assert isinstance(tensor, np.ndarray)  # type checking
     itype = np_dtype_to_tensor_dtype(tensor.dtype)
     stat = dict(
         mean=float(tensor.mean()),
@@ -948,7 +918,7 @@ class NodeCoordinates:
 
     def __init__(
         self,
-        node: Union[onnx.TensorProto, NodeProto, str],
+        node: Union[TensorProto, NodeProto, onnx.SparseTensorProto, ValueInfoProto, str],
         path: Tuple[Tuple[int, str, str], ...],
     ):
         assert isinstance(path, tuple), f"Unexpected type {type(path)} for path"
@@ -968,9 +938,7 @@ class NodeCoordinates:
 
 
 class ResultFound:
-    """
-    Class returned by :func:`enumerate_results`.
-    """
+    """Class returned by :func:`enumerate_results`."""
 
     __slots__ = ("consumer", "name", "producer")
 
@@ -1060,9 +1028,9 @@ def enumerate_results(
                     print(f"[enumerate_results] {indent}-- {r}")
                 yield r
         for i in proto.sparse_initializer:
-            if i.name in name:
+            if i.values.name in name:
                 r = ResultFound(
-                    i.name,
+                    i.values.name,
                     NodeCoordinates(i, tuple([*coordinates, (-1, "INIT", "")])),  # noqa: C409
                     None,
                 )
@@ -1165,9 +1133,9 @@ def shadowing_names(
         return shadowing_names(
             proto.node,
             verbose=verbose,
-            existing=set(i.name for i in proto.initializer)
-            | set(i.name for i in proto.sparse_initializer)
-            | set(i.name for i in proto.input if i.name),
+            existing={i.name for i in proto.initializer}
+            | {i.values.name for i in proto.sparse_initializer}
+            | {i.name for i in proto.input if i.name},
             shadow_context=set(),
             post_shadow_context=set(),
         )
@@ -1201,9 +1169,9 @@ def shadowing_names(
         for att in node.attribute:
             if att.type == AttributeProto.GRAPH:
                 g = att.g
-                shadow |= set(i.name for i in g.input) & shadow_context
-                shadow |= set(i.name for i in g.initializer) & shadow_context
-                shadow |= set(i.name for i in g.sparse_initializer) & shadow_context
+                shadow |= {i.name for i in g.input} & shadow_context
+                shadow |= {i.name for i in g.initializer} & shadow_context
+                shadow |= {i.values.name for i in g.sparse_initializer} & shadow_context
                 s, _ps, c = shadowing_names(
                     g.node, verbose=verbose, existing=existing, shadow_context=existing
                 )
@@ -1225,9 +1193,9 @@ def get_hidden_inputs(graph: onnx.GraphProto) -> Set[str]:
     """
     hidden = set()
     memo = (
-        set(i.name for i in graph.initializer)
-        | set(i.name for i in graph.sparse_initializer)
-        | set(i.name for i in graph.input)
+        {i.name for i in graph.initializer}
+        | {i.values.name for i in graph.sparse_initializer}
+        | {i.name for i in graph.input}
     )
     for node in graph.node:
         for i in node.input:
@@ -1353,7 +1321,6 @@ def make_submodel(
     Creates a model with the given list of nodes.
     It computes the minimum list of inputs needed for this model.
     The function assumes the nodes are sorted.
-    It does not handle yet subgraphs.
 
     :param nodes: list of nodes
     :param ir_version: ir version
@@ -1376,25 +1343,61 @@ def make_submodel(
                 if att.type == onnx.AttributeProto.GRAPH:
                     not_known |= get_hidden_inputs(att.g)
 
-    model = oh.make_model(
+    return oh.make_model(
         oh.make_graph(
             nodes,
             "submodel",
             [_mkv_(n, *type_rank_fn(n)) for n in sorted(not_known) if n],
-            [_mkv_(n, *type_rank_fn(n)) for n in sorted(output_names) if n],
+            [_mkv_(n, *type_rank_fn(n)) for n in output_names if n],
         ),
         ir_version=ir_version,
         opset_imports=opset_imports,
     )
-    return model
+
+
+def make_subfunction(
+    name: str,
+    nodes: List[NodeProto],
+    opset_imports: Sequence[OperatorSetIdProto],
+    output_names: List[str],
+    domain: str = "local_function",
+) -> FunctionProto:
+    """
+    Creates a function with the given list of nodes.
+    It computes the minimum list of inputs needed for this model.
+    The function assumes the nodes are sorted.
+
+    :param name: function name
+    :param nodes: list of nodes
+    :param opset_imports: opset import
+    :param output_names: desired outputs
+    :param domain: function domain
+    :return: model proto
+    """
+    not_known: Set[str] = set()
+    for node in nodes[::-1]:
+        not_known -= {o for o in node.output if o}
+        not_known |= {i for i in node.input if i}
+        if node.op_type in {"Scan", "If", "Loop"}:
+            # there are hidden inputs
+            for att in node.attribute:
+                if att.type == onnx.AttributeProto.GRAPH:
+                    not_known |= get_hidden_inputs(att.g)
+
+    return oh.make_function(
+        domain,
+        name,
+        nodes=nodes,
+        inputs=sorted(not_known),
+        outputs=output_names,
+        opset_imports=opset_imports,
+    )
 
 
 def get_tensor_shape(
     obj: Union[onnx.ValueInfoProto, onnx.TypeProto, onnx.TensorProto],
 ) -> Optional[List[Optional[Union[int, str]]]]:
-    """
-    Returns the shape if that makes sense for this object.
-    """
+    """Returns the shape if that makes sense for this object."""
     if isinstance(obj, ValueInfoProto):
         return get_tensor_shape(obj.type)
     elif not isinstance(obj, onnx.TypeProto):
@@ -1512,9 +1515,6 @@ def onnx_remove_node_unused(
         if not ({o for o in node.output if o} & marked_set):
             removed.add(ind)
 
-    if not is_function:
-        initializers = [i for i in graph.initializer if i.name in marked]
-        sparse_initializers = [i for i in graph.sparse_initializer if i.name in marked]
     new_nodes = [node for i, node in enumerate(nodes) if i not in removed]
 
     # Finally create the new graph.
@@ -1529,13 +1529,16 @@ def onnx_remove_node_unused(
             attributes=graph.attribute,
             doc_string=graph.doc_string,
         )
+
+    initializers = [i for i in graph.initializer if i.name in marked]
+    sparse_initializers = [i for i in graph.sparse_initializer if i.values.name in marked]
     new_graph = oh.make_graph(
         new_nodes,
         graph.name,
         graph.input,
         graph.output,
         initializers,
-        sparse_initializers,
+        sparse_initializer=sparse_initializers,
     )
     new_graph.value_info.extend(graph.value_info)
     return new_graph
@@ -1549,7 +1552,7 @@ def select_model_inputs_outputs(
     overwrite: Optional[Dict[str, Any]] = None,
     remove_unused: bool = True,
     verbose: int = 0,
-):
+) -> ModelProto:
     """
     Takes a model and changes its outputs.
 
@@ -1709,6 +1712,7 @@ def select_model_inputs_outputs(
     )
     if remove_unused:
         graph = onnx_remove_node_unused(graph, recursive=False)
+        assert isinstance(graph, GraphProto)  # type checking
     onnx_model = oh.make_model(graph, functions=model.functions)
     onnx_model.ir_version = model.ir_version
     onnx_model.producer_name = model.producer_name
@@ -1727,3 +1731,172 @@ def select_model_inputs_outputs(
         op_set.version = oimp.version
 
     return onnx_model
+
+
+def _find_used_names(node_list, node_indices):
+    # find all the outputs the subset of nodes produces
+    possible_outputs = set()
+    for i_node in node_indices:
+        if not node_list[i_node]:
+            continue
+        possible_outputs |= {o for o in node_list[i_node].output if o}
+    # find all requires input from the other nodes
+    set_indices = set(node_indices)
+    not_known: Set[str] = set()
+    ranges = list(range(len(node_list)))
+    for i_node in ranges[::-1]:
+        if i_node in set_indices:
+            continue
+        node = node_list[i_node]
+        if not node:
+            continue
+        not_known -= {o for o in node.output if o}
+        not_known |= {i for i in node.input if i}
+        if node.op_type in {"Scan", "If", "Loop"}:
+            # there are hidden inputs
+            for att in node.attribute:
+                if att.type == onnx.AttributeProto.GRAPH:
+                    not_known |= get_hidden_inputs(att.g)
+    # output
+    selection = possible_outputs & not_known
+    assert selection, (
+        f"No output is needed, possible_outputs={sorted(possible_outputs)}, "
+        f"not_known={sorted(not_known)}"
+    )
+    return sorted(selection)
+
+
+def check_for_non_recursivity(
+    node_list: List[Optional[NodeProto]], inputs: Sequence[str], outputs: Sequence[str]
+):
+    """
+    We finally need to check that any of this output is not required
+    by one input from the function itself, that would mean one node
+    needs an output of the function and is also required by the function:
+    it is probably missing from the initial set.
+
+
+
+    :param node_list: list of nodes
+    :param inputs: input names to consider
+    :param outputs: output names which cannot be involved in input names
+    """
+    set_inputs = set(inputs)
+    set_outputs = set(outputs)
+    for node in node_list[::-1]:
+        if not node:
+            continue
+        si = set(node.output)
+        if si & set_inputs:
+            set_inputs |= set(node.input)
+            if node.op_type in {"Scan", "If", "Loop"}:
+                # there are hidden inputs
+                for att in node.attribute:
+                    if att.type == onnx.AttributeProto.GRAPH:
+                        set_inputs |= get_hidden_inputs(att.g)
+        if set_outputs & set_inputs:
+            raise ValueError(
+                f"Results {set_outputs & set_inputs} are needed for inputs {inputs} "
+                f"but also requires {outputs} which is not allowed."
+            )
+
+
+def make_model_with_local_functions(
+    model: ModelProto,
+    regex: str = ".*[.]layers[.][0-9]+[.]forward$",
+    domain: str = "local_function",
+    metadata_key_prefix: Union[str, Tuple[str, ...]] = ("namespace", "source["),
+    verbose: int = 0,
+) -> ModelProto:
+    """
+    Selects nodes based on a regular expression, using metadata
+    ``'namespace'``. It is going to look into every value
+    matching the regular expression and partition the nodes based
+    on the unique values the regular expression finds.
+    Every set of nodes it replaced by a call to a local function.
+
+    :param model: model proto
+    :param regex: regular expression
+    :param domain: function domain
+    :param metadata_keys: list of metadata keys to consider,
+        every value is split into multiple ones.
+    :param verbose: verbosity
+    :return: model proto
+    """
+    prefix = (
+        metadata_key_prefix
+        if isinstance(metadata_key_prefix, tuple)
+        else (metadata_key_prefix,)
+    )
+    reg = re.compile(regex)
+    unique_values = set()
+    unique: Dict[str, List[int]] = {}
+    for i, node in enumerate(model.graph.node):
+        selected = False
+        for data in node.metadata_props:
+            if data.key.startswith(prefix):
+                values = re.split("[,:]", data.value)
+                for v in values:
+                    if not v:
+                        continue
+                    if reg.match(v):
+                        if v not in unique:
+                            unique[v] = []
+                        unique[v].append(i)
+                        selected = True
+                        break
+                    unique_values.add(v)
+                if selected:
+                    break
+    # sets of nodes.
+    if not unique:
+        if verbose:
+            print(f"[make_model_with_local_functions] no match in {sorted(unique_values)}")
+        return model
+
+    if verbose:
+        print(f"[make_model_with_local_functions] matched {len(unique)} partitions")
+    functions = []
+    new_nodes: List[Optional[NodeProto]] = list(model.graph.node)
+    for key, node_indices in unique.items():
+        function_name = key.strip().replace(".", "_")
+        if verbose:
+            print(
+                f"[make_model_with_local_functions] move {len(node_indices)} "
+                f"nodes in partition {function_name!r}"
+            )
+        outputs = _find_used_names(new_nodes, node_indices)
+        function_nodes = [new_nodes[i] for i in node_indices]
+        lf = make_subfunction(
+            function_name,
+            [n for n in function_nodes if n],
+            model.opset_import,
+            outputs,
+            domain=domain,
+        )
+        check_for_non_recursivity(new_nodes, lf.input, lf.output)
+        functions.append(lf)
+        maxi = max(node_indices)
+        for i in node_indices:
+            new_nodes[i] = None
+        new_nodes[maxi] = oh.make_node(lf.name, lf.input, lf.output, domain=lf.domain)
+
+    return oh.make_model(
+        oh.make_graph(
+            [n for n in new_nodes if n],
+            model.graph.name,
+            model.graph.input,
+            model.graph.output,
+            model.graph.initializer,
+            doc_string=model.graph.doc_string,
+            value_info=model.graph.value_info,
+            sparse_initializer=model.graph.sparse_initializer,
+        ),
+        ir_version=model.ir_version,
+        opset_imports=(
+            model.opset_import
+            if domain in {d.domain for d in model.opset_import}
+            else [*model.opset_import, oh.make_opsetid(domain, 1)]
+        ),
+        functions=[*model.functions, *functions],
+    )

onnx_diagnostic/helpers/ort_session.py

@@ -14,7 +14,6 @@ from .onnx_helper import (
 )
 from .torch_helper import torch_dtype_to_onnx_dtype
 
-
 DEVICES = {-1: ORTC.OrtDevice(ORTC.OrtDevice.cpu(), ORTC.OrtDevice.default_memory(), 0)}
 TensorLike = Union[np.ndarray, torch.Tensor]
 

onnx_diagnostic/helpers/torch_helper.py

@@ -19,12 +19,7 @@ from .cache_helper import (
     CacheKeyValue,
 )
 from .mini_onnx_builder import create_onnx_model_from_input_tensors
-from .onnx_helper import (
-    to_array_extended,
-    tensor_dtype_to_np_dtype,
-    _STORAGE_TYPE,
-    onnx_dtype_name,
-)
+from .onnx_helper import to_array_extended, tensor_dtype_to_np_dtype, onnx_dtype_name
 
 
 def proto_from_tensor(
@@ -84,13 +79,17 @@ def proto_from_tensor(
         byte_data = (ctypes.c_ubyte * numel * element_size).from_address(np_arr.data_ptr())
         tensor.raw_data = bytes(byte_data)
         if sys.byteorder == "big":
-            np_dtype = _STORAGE_TYPE[tensor.data_type]  # type: ignore
-            np.byteswap(np.frombuffer(tensor.raw_data, dtype=np_dtype), inplace=True)  # type: ignore
+            storage_type = {
+                onnx.TensorProto.FLOAT16: np.int16,
+                onnx.TensorProto.BFLOAT16: np.int16,
+            }
+            np_dtype = storage_type[tensor.data_type]  # type: ignore
+            np.frombuffer(tensor.raw_data, dtype=np_dtype).byteswap(inplace=True)  # type: ignore
     else:
         tensor.raw_data = np_arr.tobytes()
         if sys.byteorder == "big":
             np_dtype = tensor_dtype_to_np_dtype(tensor.data_type)
-            np.byteswap(np.frombuffer(tensor.raw_data, dtype=np_dtype), inplace=True)
+            np.frombuffer(tensor.raw_data, dtype=np_dtype).byteswap(inplace=True)
     return tensor