Add to_tensor

xadupre · xadupre · commit 73fba13af593 · 2025-05-13T10:23:50.000+02:00
diff --git a/_doc/api/helpers/index.rst b/_doc/api/helpers/index.rst
@@ -16,7 +16,7 @@ onnx_diagnostic.helpers
     onnx_helper
     ort_session
     rt_helper
-    torch_test_helper
+    torch_helper
 
 .. autofunction:: onnx_diagnostic.helpers.flatten_object
 
diff --git a/_doc/api/helpers/torch_helper.rst b/_doc/api/helpers/torch_helper.rst
@@ -0,0 +1,7 @@
+
+onnx_diagnostic.helpers.torch_helper
+====================================
+
+.. automodule:: onnx_diagnostic.helpers.torch_helper
+    :members:
+    :no-undoc-members:
diff --git a/_doc/api/helpers/torch_test_helper.rst b/_doc/api/helpers/torch_test_helper.rst
diff --git a/_doc/examples/plot_export_tiny_llm.py b/_doc/examples/plot_export_tiny_llm.py
@@ -31,7 +31,7 @@
 import transformers
 from onnx_diagnostic import doc
 from onnx_diagnostic.helpers import string_type
-from onnx_diagnostic.helpers.torch_test_helper import steal_forward
+from onnx_diagnostic.helpers.torch_helper import steal_forward
 from onnx_diagnostic.torch_models.llms import get_tiny_llm
 
 
@@ -77,7 +77,7 @@ def _forward_(*args, _f=None, **kwargs):
 model.forward = keep_model_forward
 
 # %%
-# Another syntax with :func:`onnx_diagnostic.helpers.torch_test_helper.steal_forward`.
+# Another syntax with :func:`onnx_diagnostic.helpers.torch_helper.steal_forward`.
 
 with steal_forward(model):
     model.generate(inputs, max_length=50, temperature=1, top_k=50, top_p=0.95, do_sample=True)
diff --git a/_unittests/ut_export/test_jit.py b/_unittests/ut_export/test_jit.py
@@ -8,7 +8,7 @@
     requires_onnxscript,
 )
 from onnx_diagnostic.reference import ExtendedReferenceEvaluator
-from onnx_diagnostic.helpers.torch_test_helper import is_torchdynamo_exporting
+from onnx_diagnostic.helpers.torch_helper import is_torchdynamo_exporting
 
 try:
     from experimental_experiment.torch_interpreter import to_onnx
diff --git a/_unittests/ut_helpers/test_torch_helper.py b/_unittests/ut_helpers/test_torch_helper.py
@@ -1,11 +1,12 @@
 import unittest
+import numpy as np
 import ml_dtypes
 import onnx
 import torch
 import transformers
 from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout
 from onnx_diagnostic.helpers import max_diff, string_type
-from onnx_diagnostic.helpers.torch_test_helper import (
+from onnx_diagnostic.helpers.torch_helper import (
     dummy_llm,
     to_numpy,
     is_torchdynamo_exporting,
@@ -24,6 +25,8 @@
     make_sliding_window_cache,
 )
 from onnx_diagnostic.helpers.mini_onnx_builder import create_input_tensors_from_onnx_model
+from onnx_diagnostic.helpers.onnx_helper import from_array_extended, to_array_extended
+from onnx_diagnostic.helpers.torch_helper import to_tensor
 
 TFLOAT = onnx.TensorProto.FLOAT
 
@@ -205,7 +208,7 @@ def forward(self, x, y):
             else:
                 print("output", k, v)
         print(string_type(restored, with_shape=True))
-        l1, l2 = 183, 192
+        l1, l2 = 186, 195
         self.assertEqual(
             [
                 (f"-Model-{l2}", 0, "I"),
@@ -344,6 +347,35 @@ def forward(self, x, y=None):
             stat,
         )
 
+    def test_to_tensor(self):
+        for dtype in [
+            np.int8,
+            np.uint8,
+            np.int16,
+            np.uint16,
+            np.int32,
+            np.uint32,
+            np.int64,
+            np.uint64,
+            np.float16,
+            np.float32,
+            np.float64,
+        ]:
+            with self.subTest(dtype=dtype):
+                a = np.random.rand(4, 5).astype(dtype)
+                proto = from_array_extended(a)
+                b = to_array_extended(proto)
+                self.assertEqualArray(a, b)
+                c = to_tensor(proto)
+                self.assertEqualArray(a, c)
+
+        for dtype in [torch.bfloat16]:
+            with self.subTest(dtype=dtype):
+                a = torch.rand((4, 5), dtype=dtype)
+                proto = from_array_extended(a)
+                c = to_tensor(proto)
+                self.assertEqualArray(a, c)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)
diff --git a/_unittests/ut_tasks/try_tasks.py b/_unittests/ut_tasks/try_tasks.py
@@ -1,7 +1,7 @@
 import unittest
 from onnx_diagnostic.ext_test_case import ExtTestCase, never_test
 from onnx_diagnostic.helpers import string_type
-from onnx_diagnostic.helpers.torch_test_helper import steal_forward
+from onnx_diagnostic.helpers.torch_helper import steal_forward
 
 
 class TestHuggingFaceHubModel(ExtTestCase):
diff --git a/_unittests/ut_torch_export_patches/test_patch_expressions.py b/_unittests/ut_torch_export_patches/test_patch_expressions.py
@@ -7,7 +7,7 @@
     patched_selector,
     patched_float_arange,
 )
-from onnx_diagnostic.helpers.torch_test_helper import fake_torchdynamo_exporting
+from onnx_diagnostic.helpers.torch_helper import fake_torchdynamo_exporting
 
 
 class TestOnnxExportErrors(ExtTestCase):
diff --git a/_unittests/ut_torch_export_patches/test_patch_loops.py b/_unittests/ut_torch_export_patches/test_patch_loops.py
@@ -1,7 +1,7 @@
 import unittest
 import torch
 from onnx_diagnostic.ext_test_case import ExtTestCase, requires_torch
-from onnx_diagnostic.helpers.torch_test_helper import (
+from onnx_diagnostic.helpers.torch_helper import (
     is_torchdynamo_exporting,
     fake_torchdynamo_exporting,
 )
diff --git a/_unittests/ut_torch_export_patches/test_patch_serialization.py b/_unittests/ut_torch_export_patches/test_patch_serialization.py
@@ -11,7 +11,7 @@
 from onnx_diagnostic.torch_export_patches.onnx_export_errors import (
     torch_export_patches,
 )
-from onnx_diagnostic.helpers.torch_test_helper import torch_deepcopy
+from onnx_diagnostic.helpers.torch_helper import torch_deepcopy
 
 
 class TestPatchSerialization(ExtTestCase):
diff --git a/onnx_diagnostic/export/validate.py b/onnx_diagnostic/export/validate.py
@@ -4,7 +4,7 @@
 from typing import Any, Dict, List, Optional, Tuple, Union
 import torch
 from ..helpers import string_type, max_diff, string_diff
-from ..helpers.torch_test_helper import torch_deepcopy
+from ..helpers.torch_helper import torch_deepcopy
 from .dynamic_shapes import CoupleInputsDynamicShapes
 
 
diff --git a/onnx_diagnostic/ext_test_case.py b/onnx_diagnostic/ext_test_case.py
@@ -871,7 +871,7 @@ def assertEqualArray(
                 raise AssertionError("\n".join(rows))  # noqa: B904
             return
 
-        from .helpers.torch_test_helper import to_numpy
+        from .helpers.torch_helper import to_numpy
 
         if hasattr(expected, "detach"):
             expected = to_numpy(expected.detach().cpu())
diff --git a/onnx_diagnostic/helpers/helper.py b/onnx_diagnostic/helpers/helper.py
@@ -439,7 +439,7 @@ def string_type(
             print(f"[string_type] F2:{type(obj)}")
         return f"{prefix}F{i}s{'x'.join(map(str, obj.shape))}"
     if isinstance(obj, torch.Tensor):
-        from .onnx_helper import torch_dtype_to_onnx_dtype
+        from .torch_helper import torch_dtype_to_onnx_dtype
 
         if with_min_max:
             s = string_type(obj, with_shape=with_shape, with_device=with_device)
diff --git a/onnx_diagnostic/helpers/onnx_helper.py b/onnx_diagnostic/helpers/onnx_helper.py
@@ -1,4 +1,3 @@
-import ctypes
 import functools
 import json
 import os
@@ -518,76 +517,6 @@ def from_array_ml_dtypes(arr: npt.ArrayLike, name: Optional[str] = None) -> Tens
 }
 
 
-def proto_from_tensor(
-    arr: "torch.Tensor",  # noqa: F821
-    name: Optional[str] = None,
-    verbose: int = 0,
-) -> TensorProto:
-    """
-    Converts a torch Tensor into a TensorProto.
-
-    :param arr: tensor
-    :param verbose: display the type and shape
-    :return: a TensorProto
-    """
-    import torch
-
-    if not isinstance(arr, torch.Tensor):
-        raise TypeError(f"Unexpected type {type(arr)}.")
-    if arr.is_sparse:
-        raise NotImplementedError(
-            f"Sparse tensor is not supported yet but initializer {name!r} is."
-        )
-
-    # arr.contiguous() is slow after a transpose, maybe there is a way to optimize this.
-    if arr.is_contiguous():
-        arr_cpu = arr.cpu()
-    else:
-        arr_cpu = arr.contiguous().cpu()
-
-    numel = torch.numel(arr_cpu)
-    element_size = arr_cpu.element_size()
-
-    if arr_cpu.dtype in {torch.bfloat16}:
-        np_arr = arr_cpu
-    elif arr_cpu.data_ptr() == arr.data_ptr():
-        copy = arr_cpu.clone().detach().requires_grad_(False)
-        assert (
-            arr_cpu.data_ptr() == 0 or arr_cpu.data_ptr() != copy.data_ptr()
-        ), f"Pointers are not null and different {arr_cpu.data_ptr()} != {copy.data_ptr()}"
-        np_arr = np.from_dlpack(copy)
-    else:
-        np_arr = np.from_dlpack(arr_cpu.detach())
-
-    tensor = TensorProto()
-    tensor.dims.extend(arr_cpu.shape)
-    if name:
-        tensor.name = name
-    itype = torch_dtype_to_onnx_dtype(arr_cpu.dtype)
-    assert not hasattr(TensorProto, "INT4") or itype not in {
-        TensorProto.INT4,
-        TensorProto.UINT4,
-    }, f"Type {arr.dtype} is not supported yet for name={name!r}"
-    tensor.data_type = itype
-
-    if verbose > 1 and numel > 100:
-        print(f"[proto_from_array] {tensor.data_type}[{arr_cpu.shape}]")
-
-    if isinstance(np_arr, torch.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
-    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)
-
-    return tensor
-
-
 def from_array_extended(tensor: npt.ArrayLike, name: Optional[str] = None) -> TensorProto:
     """
     Converts an array into a :class:`onnx.TensorProto`.
@@ -596,11 +525,13 @@ def from_array_extended(tensor: npt.ArrayLike, name: Optional[str] = None) -> Te
     :param name: name
     :return: TensorProto
     """
-    try:
+    if not isinstance(tensor, np.ndarray):
         import torch
-    except ImportError:
-        torch = None
-    if torch is not None and isinstance(tensor, torch.Tensor):
+        from .torch_helper import proto_from_tensor
+
+        assert isinstance(
+            tensor, torch.Tensor
+        ), f"Unable to convert type {type(tensor)} into TensorProto."
         return proto_from_tensor(tensor, name=name)
 
     from onnx.reference.ops.op_cast import (
@@ -657,50 +588,6 @@ def to_array_extended(proto: TensorProto) -> npt.ArrayLike:
     return arr
 
 
-def onnx_dtype_to_torch_dtype(itype: int) -> "torch.dtype":  # noqa: F821
-    """
-    Converts an onnx type into a torch dtype.
-
-    :param to: onnx dtype
-    :return: torch dtype
-    """
-    import torch
-
-    if itype == TensorProto.FLOAT:
-        return torch.float32
-    if itype == TensorProto.FLOAT16:
-        return torch.float16
-    if itype == TensorProto.BFLOAT16:
-        return torch.bfloat16
-    if itype == TensorProto.DOUBLE:
-        return torch.float64
-    if itype == TensorProto.INT32:
-        return torch.int32
-    if itype == TensorProto.INT64:
-        return torch.int64
-    if itype == TensorProto.UINT32:
-        return torch.uint32
-    if itype == TensorProto.UINT64:
-        return torch.uint64
-    if itype == TensorProto.BOOL:
-        return torch.bool
-    if itype == TensorProto.INT16:
-        return torch.int16
-    if itype == TensorProto.UINT16:
-        return torch.uint16
-    if itype == TensorProto.INT8:
-        return torch.int16
-    if itype == TensorProto.UINT8:
-        return torch.uint16
-    if itype == TensorProto.COMPLEX64:
-        return torch.complex64
-    if itype == TensorProto.COMPLEX128:
-        return torch.complex128
-    raise NotImplementedError(
-        f"Unable to convert onnx type {onnx_dtype_name(itype)} to torch.type."
-    )
-
-
 def onnx_dtype_to_np_dtype(itype: int) -> Any:
     """
     Converts an onnx type into a to numpy dtype.
@@ -746,52 +633,6 @@ def onnx_dtype_to_np_dtype(itype: int) -> Any:
     )
 
 
-def torch_dtype_to_onnx_dtype(to: "torch.dtype") -> int:  # noqa: F821
-    """
-    Converts a torch dtype into a onnx element type.
-
-    :param to: torch dtype
-    :return: onnx type
-    """
-    import torch
-
-    if to == torch.float32:
-        return TensorProto.FLOAT
-    if to == torch.float16:
-        return TensorProto.FLOAT16
-    if to == torch.bfloat16:
-        return TensorProto.BFLOAT16
-    if to == torch.float64:
-        return TensorProto.DOUBLE
-    if to == torch.int64:
-        return TensorProto.INT64
-    if to == torch.int32:
-        return TensorProto.INT32
-    if to == torch.uint64:
-        return TensorProto.UINT64
-    if to == torch.uint32:
-        return TensorProto.UINT32
-    if to == torch.bool:
-        return TensorProto.BOOL
-    if to == torch.SymInt:
-        return TensorProto.INT64
-    if to == torch.int16:
-        return TensorProto.INT16
-    if to == torch.uint16:
-        return TensorProto.UINT16
-    if to == torch.int8:
-        return TensorProto.INT8
-    if to == torch.uint8:
-        return TensorProto.UINT8
-    if to == torch.SymFloat:
-        return TensorProto.FLOAT
-    if to == torch.complex64:
-        return TensorProto.COMPLEX64
-    if to == torch.complex128:
-        return TensorProto.COMPLEX128
-    raise NotImplementedError(f"Unable to convert torch dtype {to!r} to onnx dtype.")
-
-
 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.
@@ -803,6 +644,8 @@ def dtype_to_tensor_dtype(dt: Union[np.dtype, "torch.dtype"]) -> int:  # noqa: F
         return np_dtype_to_tensor_dtype(dt)
     except (KeyError, TypeError, ValueError):
         pass
+    from .torch_helper import torch_dtype_to_onnx_dtype
+
     return torch_dtype_to_onnx_dtype(dt)
 
 
diff --git a/onnx_diagnostic/helpers/ort_session.py b/onnx_diagnostic/helpers/ort_session.py
diff --git a/onnx_diagnostic/helpers/torch_helper.py b/onnx_diagnostic/helpers/torch_helper.py
diff --git a/onnx_diagnostic/torch_export_patches/patch_expressions.py b/onnx_diagnostic/torch_export_patches/patch_expressions.py
diff --git a/onnx_diagnostic/torch_export_patches/patches/patch_transformers.py b/onnx_diagnostic/torch_export_patches/patches/patch_transformers.py
diff --git a/onnx_diagnostic/torch_models/test_helper.py b/onnx_diagnostic/torch_models/test_helper.py

Original file line number	Diff line number	Diff line change
`@@ -8,7 +8,7 @@`
`8`	`8`	`requires_onnxscript,`
`9`	`9`	`)`
`10`	`10`	`from onnx_diagnostic.reference import ExtendedReferenceEvaluator`
`11`		`-from onnx_diagnostic.helpers.torch_test_helper import is_torchdynamo_exporting`
	`11`	`+from onnx_diagnostic.helpers.torch_helper import is_torchdynamo_exporting`
`12`	`12`
`13`	`13`	`try:`
`14`	`14`	`from experimental_experiment.torch_interpreter import to_onnx`
Original file line number	Diff line number	Diff line change
`@@ -7,7 +7,7 @@`
`7`	`7`	`patched_selector,`
`8`	`8`	`patched_float_arange,`
`9`	`9`	`)`
`10`		`-from onnx_diagnostic.helpers.torch_test_helper import fake_torchdynamo_exporting`
	`10`	`+from onnx_diagnostic.helpers.torch_helper import fake_torchdynamo_exporting`
`11`	`11`
`12`	`12`
`13`	`13`	`class TestOnnxExportErrors(ExtTestCase):`
Original file line number	Diff line number	Diff line change
`@@ -1,7 +1,7 @@`
`1`	`1`	`import unittest`
`2`	`2`	`import torch`
`3`	`3`	`from onnx_diagnostic.ext_test_case import ExtTestCase, requires_torch`
`4`		`-from onnx_diagnostic.helpers.torch_test_helper import (`
	`4`	`+from onnx_diagnostic.helpers.torch_helper import (`
`5`	`5`	`is_torchdynamo_exporting,`
`6`	`6`	`fake_torchdynamo_exporting,`
`7`	`7`	`)`
Original file line number	Diff line number	Diff line change
`@@ -11,7 +11,7 @@`
`11`	`11`	`from onnx_diagnostic.torch_export_patches.onnx_export_errors import (`
`12`	`12`	`torch_export_patches,`
`13`	`13`	`)`
`14`		`-from onnx_diagnostic.helpers.torch_test_helper import torch_deepcopy`
	`14`	`+from onnx_diagnostic.helpers.torch_helper import torch_deepcopy`
`15`	`15`
`16`	`16`
`17`	`17`	`class TestPatchSerialization(ExtTestCase):`