Repository: onnx/onnx
Pull request: Improve mapping and add more tests for make_tensor
**Describe your changes**
- Refactor mapping
- Deprecate TENSOR_TYPE_TO_STORAGE_TENSOR_TYPE
- Use functions instead of more maps to prevent confusions and save binary size
- Add more tests (test_make_tensor) for more tensor types 

**Motivation and Context**
- Relevant to https://github.com/onnx/onnx/issues/4261. Current mapping is quite confusing and there are some duplicate implementation for getting field by a tensor type.
- Now there are missing some tensor types for test_make_tensor w/ and w/o raw.

Problem statement:
provide a nicer API for accessing the proper member of TensorProto
Currently code like this is repeated several times:

```py
field = mapping.STORAGE_TENSOR_TYPE_TO_FIELD[
    mapping.TENSOR_TYPE_TO_STORAGE_TENSOR_TYPE[data_type]]
getattr(tensor, field)
```

This is repetitive and can be encapsulated in a helper function.
Also the name "storage tensor type" is misleading and has led to at least one bug.

We should:
1. Add a function that does all this
2. Deprecate `mapping.STORAGE_TENSOR_TYPE_TO_FIELD` and `mapping.TENSOR_TYPE_TO_STORAGE_TENSOR_TYPE`. Ideally log a warning when they're accessed.
3. Eventually, then in the next release, remove those fields entirely.

Hints from issue discussion:

Commit message: fix mypy typecheck failures

Signed-off-by: Chun-Wei Chen <jacky82226@gmail.com>

Patch to review:
diff --git a/onnx/helper.py b/onnx/helper.py
index 21396aa3075..3c890487ec3 100644
--- a/onnx/helper.py
+++ b/onnx/helper.py
@@ -362,8 +362,7 @@ def make_tensor(
             vals = np.array(vals).astype(np_dtype).view(dtype=np.uint16).flatten().tolist()
         elif data_type == TensorProto.BFLOAT16:
             vals = list(map(float32_to_bfloat16, np.array(vals).astype(np_dtype).flatten().tolist()))
-        field = mapping.STORAGE_TENSOR_TYPE_TO_FIELD[
-            mapping.TENSOR_TYPE_TO_STORAGE_TENSOR_TYPE[data_type]]
+        field = mapping.to_field(data_type)
         getattr(tensor, field).extend(vals)
     tensor.dims.extend(dims)
     return tensor
diff --git a/onnx/mapping.py b/onnx/mapping.py
index be5b688fd74..96de9161d6f 100644
--- a/onnx/mapping.py
+++ b/onnx/mapping.py
@@ -2,52 +2,42 @@
 
 from onnx import TensorProto, SequenceProto, OptionalProto
 import numpy as np  # type: ignore
+from typing import Any
+import warnings
 
-# This map is used for converting TensorProto values into Numpy arrays
-TENSOR_TYPE_TO_NP_TYPE = {
-    int(TensorProto.FLOAT): np.dtype('float32'),
-    int(TensorProto.UINT8): np.dtype('uint8'),
-    int(TensorProto.INT8): np.dtype('int8'),
-    int(TensorProto.UINT16): np.dtype('uint16'),
-    int(TensorProto.INT16): np.dtype('int16'),
-    int(TensorProto.INT32): np.dtype('int32'),
-    int(TensorProto.INT64): np.dtype('int64'),
-    int(TensorProto.BOOL): np.dtype('bool'),
-    int(TensorProto.FLOAT16): np.dtype('float16'),
-    int(TensorProto.BFLOAT16): np.dtype('float32'),  # Native numpy does not support bfloat16 so now use float32 for bf16 values
-    int(TensorProto.DOUBLE): np.dtype('float64'),
-    int(TensorProto.COMPLEX64): np.dtype('complex64'),
-    int(TensorProto.COMPLEX128): np.dtype('complex128'),
-    int(TensorProto.UINT32): np.dtype('uint32'),
-    int(TensorProto.UINT64): np.dtype('uint64'),
-    int(TensorProto.STRING): np.dtype('object')
+# tesnor_type: (numpy type, storage type, string name)
+TENSOR_TYPE_MAP = {
+    int(TensorProto.FLOAT): (np.dtype('float32'), int(TensorProto.FLOAT), 'TensorProto.FLOAT'),
+    int(TensorProto.UINT8): (np.dtype('uint8'), int(TensorProto.INT32), 'TensorProto.UINT8'),
+    int(TensorProto.INT8): (np.dtype('int8'), int(TensorProto.INT32), 'TensorProto.INT8'),
+    int(TensorProto.UINT16): (np.dtype('uint16'), int(TensorProto.INT32), 'TensorProto.UINT16'),
+    int(TensorProto.INT16): (np.dtype('int16'), int(TensorProto.INT32), 'TensorProto.INT16'),
+    int(TensorProto.INT32): (np.dtype('int32'), int(TensorProto.INT32), 'TensorProto.INT32'),
+    int(TensorProto.INT64): (np.dtype('int64'), int(TensorProto.INT64), 'TensorProto.INT64'),
+    int(TensorProto.BOOL): (np.dtype('bool'), int(TensorProto.INT32), 'TensorProto.BOOL'),
+    int(TensorProto.FLOAT16): (np.dtype('float16'), int(TensorProto.UINT16), 'TensorProto.FLOAT16'),
+    # Native numpy does not support bfloat16 so now use float32 for bf16 values
+    # TODO ONNX should dirtectly use bfloat16 for bf16 values after numpy has supported bfloat16 type
+    int(TensorProto.BFLOAT16): (np.dtype('float32'), int(TensorProto.UINT16), 'TensorProto.BFLOAT16'),
+    int(TensorProto.DOUBLE): (np.dtype('float64'), int(TensorProto.DOUBLE), 'TensorProto.DOUBLE'),
+    int(TensorProto.COMPLEX64): (np.dtype('complex64'), int(TensorProto.FLOAT), 'TensorProto.COMPLEX64'),
+    int(TensorProto.COMPLEX128): (np.dtype('complex128'), int(TensorProto.DOUBLE), 'TensorProto.COMPLEX128'),
+    int(TensorProto.UINT32): (np.dtype('uint32'), int(TensorProto.UINT32), 'TensorProto.UINT32'),
+    int(TensorProto.UINT64): (np.dtype('uint64'), int(TensorProto.UINT64), 'TensorProto.UINT64'),
+    int(TensorProto.STRING): (np.dtype('object'), int(TensorProto.STRING), 'TensorProto.STRING'),
 }
 
+# This map is used for converting TensorProto values into numpy arrays
+TENSOR_TYPE_TO_NP_TYPE = {tensor_type: value[0] for tensor_type, value in TENSOR_TYPE_MAP.items()}
+# This is only used to get keys into STORAGE_TENSOR_TYPE_TO_FIELD.
+# TODO(https://github.com/onnx/onnx/issues/4261): Remove this.
+warnings.warn(str("`mapping.TENSOR_TYPE_TO_STORAGE_TENSOR_TYPE` is deprecated in ONNX 1.13 and will be removed in next release. To silence this warning, please use `to_sotrage_tensor_type` instead."), DeprecationWarning, stacklevel=2)
+TENSOR_TYPE_TO_STORAGE_TENSOR_TYPE = {tensor_type: value[1] for tensor_type, value in TENSOR_TYPE_MAP.items()}
+
 # Currently native numpy does not support bfloat16 so TensorProto.BFLOAT16 is ignored for now
 # Numpy float32 array is only reversed to TensorProto.FLOAT
 NP_TYPE_TO_TENSOR_TYPE = {v: k for k, v in TENSOR_TYPE_TO_NP_TYPE.items() if k != TensorProto.BFLOAT16}
 
-# This is only used to get keys into STORAGE_TENSOR_TYPE_TO_FIELD.
-# TODO(https://github.com/onnx/onnx/issues/4261): Remove this.
-TENSOR_TYPE_TO_STORAGE_TENSOR_TYPE = {
-    int(TensorProto.FLOAT): int(TensorProto.FLOAT),
-    int(TensorProto.UINT8): int(TensorProto.INT32),
-    int(TensorProto.INT8): int(TensorProto.INT32),
-    int(TensorProto.UINT16): int(TensorProto.INT32),
-    int(TensorProto.INT16): int(TensorProto.INT32),
-    int(TensorProto.INT32): int(TensorProto.INT32),
-    int(TensorProto.INT64): int(TensorProto.INT64),
-    int(TensorProto.BOOL): int(TensorProto.INT32),
-    int(TensorProto.FLOAT16): int(TensorProto.UINT16),
-    int(TensorProto.BFLOAT16): int(TensorProto.UINT16),
-    int(TensorProto.DOUBLE): int(TensorProto.DOUBLE),
-    int(TensorProto.COMPLEX64): int(TensorProto.FLOAT),
-    int(TensorProto.COMPLEX128): int(TensorProto.DOUBLE),
-    int(TensorProto.UINT32): int(TensorProto.UINT32),
-    int(TensorProto.UINT64): int(TensorProto.UINT64),
-    int(TensorProto.STRING): int(TensorProto.STRING),
-}
-
 STORAGE_TENSOR_TYPE_TO_FIELD = {
     int(TensorProto.FLOAT): 'float_data',
     int(TensorProto.INT32): 'int32_data',
@@ -62,6 +52,7 @@
     int(TensorProto.BOOL): 'int32_data',
 }
 
+
 STORAGE_ELEMENT_TYPE_TO_FIELD = {
     int(SequenceProto.TENSOR): 'tensor_values',
     int(SequenceProto.SPARSE_TENSOR): 'sparse_tensor_values',
@@ -77,3 +68,24 @@
     int(OptionalProto.MAP): 'map_value',
     int(OptionalProto.OPTIONAL): 'optional_value'
 }
+
+
+def to_np_type(tensor_type: str) -> Any:
+    return TENSOR_TYPE_MAP[tensor_type][0]
+
+
+def to_sotrage_tensor_type(tensor_type: str) -> int:
+    return TENSOR_TYPE_MAP[tensor_type][1]
+
+
+def to_string(tensor_type: str) -> str:
+    return TENSOR_TYPE_MAP[tensor_type][2]
+
+
+def to_storage_numpy_type(tensor_type: str) -> int:
+    return to_np_type(to_sotrage_tensor_type(tensor_type))
+
+
+# This map is used to get storage field for certain tensor type
+def to_field(tensor_type: str) -> str:
+    return STORAGE_TENSOR_TYPE_TO_FIELD[TENSOR_TYPE_MAP[tensor_type][1]]
diff --git a/onnx/numpy_helper.py b/onnx/numpy_helper.py
index d234468dc23..0d379bcca33 100644
--- a/onnx/numpy_helper.py
+++ b/onnx/numpy_helper.py
@@ -37,9 +37,8 @@ def to_array(tensor: TensorProto, base_dir: str = "") -> np.ndarray:
 
     tensor_dtype = tensor.data_type
     np_dtype = mapping.TENSOR_TYPE_TO_NP_TYPE[tensor_dtype]
-    storage_type = mapping.TENSOR_TYPE_TO_STORAGE_TENSOR_TYPE[tensor_dtype]
-    storage_np_dtype = mapping.TENSOR_TYPE_TO_NP_TYPE[storage_type]
-    storage_field = mapping.STORAGE_TENSOR_TYPE_TO_FIELD[storage_type]
+    storage_np_dtype = mapping.to_storage_numpy_type(tensor_dtype)
+    storage_field = mapping.to_field(tensor_dtype)
     dims = tensor.dims
 
     if tensor.data_type == TensorProto.STRING:
diff --git a/onnx/test/helper_test.py b/onnx/test/helper_test.py
index a0fcacad36f..5f0652c5afd 100644
--- a/onnx/test/helper_test.py
+++ b/onnx/test/helper_test.py
@@ -7,6 +7,8 @@
 
 from onnx import helper, defs, numpy_helper, checker
 from onnx import AttributeProto, TensorProto, GraphProto, ModelProto, OptionalProto, TypeProto, SequenceProto
+from onnx.mapping import to_string, TENSOR_TYPE_MAP
+import pytest  # type: ignore
 from typing import Any, List, Tuple
 
 import unittest
@@ -360,28 +362,7 @@ def test(opset_versions: List[Tuple[str, int]], ir_version: int) -> None:
 
 class TestHelperTensorFunctions(unittest.TestCase):
 
-    def test_make_tensor(self) -> None:
-        np_array = np.random.randn(2, 3).astype(np.float32)
-
-        tensor = helper.make_tensor(
-            name='test',
-            data_type=TensorProto.FLOAT,
-            dims=(2, 3),
-            vals=np_array
-        )
-        self.assertEqual(tensor.name, 'test')
-        np.testing.assert_equal(np_array, numpy_helper.to_array(tensor))
-
-        # use raw_data field to store the data
-        tensor = helper.make_tensor(
-            name='test',
-            data_type=TensorProto.FLOAT,
-            dims=(2, 3),
-            vals=np_array.tobytes(),
-            raw=True,
-        )
-        np.testing.assert_equal(np_array, numpy_helper.to_array(tensor))
-
+    def test_make_string_tensor(self) -> None:
         string_list = list(s.encode('utf-8') for s in ['Amy', 'Billy', 'Cindy', 'David'])
         tensor = helper.make_tensor(
             name='test',
@@ -392,53 +373,6 @@ def test_make_tensor(self) -> None:
         )
         self.assertEqual(string_list, list(tensor.string_data))
 
-    def test_make_int8_tensor(self) -> None:
-        np_array = np.random.randn(2, 3).astype(np.int8)
-
-        tensor = helper.make_tensor(
-            name='test',
-            data_type=TensorProto.INT8,
-            dims=(2, 3),
-            vals=np_array
-        )
-        self.assertEqual(tensor.name, 'test')
-        np.testing.assert_equal(np_array, numpy_helper.to_array(tensor))
-
-        # use raw_data field to store the data
-        tensor = helper.make_tensor(
-            name='test',
-            data_type=TensorProto.INT8,
-            dims=(2, 3),
-            vals=np_array.tobytes(),
-            raw=True,
-        )
-        np.testing.assert_equal(np_array, numpy_helper.to_array(tensor))
-
-    def test_make_float16_tensor(self) -> None:
-        np_array = np.random.randn(2, 3).astype(np.float16)
-
-        tensor = helper.make_tensor(
-            name='test',
-            data_type=TensorProto.FLOAT16,
-            dims=np_array.shape,
-            vals=np_array
-        )
-        self.assertEqual(tensor.name, 'test')
-        np.testing.assert_equal(np_array, numpy_helper.to_array(tensor))
-
-    def test_make_float16_tensor_raw(self) -> None:
-        np_array = np.random.randn(2, 3).astype(np.float16)
-
-        tensor = helper.make_tensor(
-            name='test',
-            data_type=TensorProto.FLOAT16,
-            dims=np_array.shape,
-            vals=np_array.view(dtype=np.uint16).flatten().tobytes(),
-            raw=True
-        )
-        self.assertEqual(tensor.name, 'test')
-        np.testing.assert_equal(np_array, numpy_helper.to_array(tensor))
-
     def test_make_bfloat16_tensor(self) -> None:
         # numpy doesn't support bf16, so we have to compute the correct result manually
         np_array = np.array([[1.0, 2.0], [3.0, 4.0], [0.099853515625, 0.099365234375], [0.0998535081744, 0.1], [np.nan, np.inf]],
@@ -660,5 +594,37 @@ def test_initializer_no_matching_graph_input(self) -> None:
   %Y_Initializer[FLOAT, 1]''' in graph_str, graph_str)
 
 
+@pytest.mark.parametrize("tensor_type",
+    [t for t in TENSOR_TYPE_MAP if t not in {TensorProto.BFLOAT16, TensorProto.STRING, TensorProto.COMPLEX64, TensorProto.COMPLEX128}],
+    ids=lambda tensor_type: to_string(tensor_type)
+)
+def test_make_tensor_vals(tensor_type: str) -> None:
+    np_array = np.random.randn(2, 3).astype(TENSOR_TYPE_MAP[tensor_type])
+    tensor = helper.make_tensor(
+        name='test',
+        data_type=tensor_type,
+        dims=np_array.shape,
+        vals=np_array
+    )
+    np.testing.assert_equal(np_array, numpy_helper.to_array(tensor))
+
+
+@pytest.mark.parametrize("tensor_type",
+    [t for t in TENSOR_TYPE_MAP if t not in {TensorProto.BFLOAT16, TensorProto.STRING}],
+    ids=lambda tensor_type: to_string(tensor_type)
+)
+def test_make_tensor_raw(tensor_type: str) -> None:
+    np_array = np.random.randn(2, 3).astype(TENSOR_TYPE_MAP[tensor_type])
+    tensor = helper.make_tensor(
+        name='test',
+        data_type=tensor_type,
+        dims=np_array.shape,
+        vals=np_array.tobytes(),
+        raw=True,
+    )
+    np.testing.assert_equal(np_array, numpy_helper.to_array(tensor))
+
+
 if __name__ == '__main__':
     unittest.main()
+    pytest.main([__file__])
diff --git a/onnx/test/hub_test.py b/onnx/test/hub_test.py
index 768d1bc23fb..80cae9ee47a 100644
--- a/onnx/test/hub_test.py
+++ b/onnx/test/hub_test.py
@@ -81,3 +81,7 @@ def test_get_model_info(self) -> None:
         hub.get_model_info("mnist", self.repo, opset=8)
         hub.get_model_info("mnist", self.repo)
         self.assertRaises(AssertionError, lambda: hub.get_model_info("mnist", self.repo, opset=-1))
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/onnx/test/parser_test.py b/onnx/test/parser_test.py
index d76c84a7d49..00d1a8175f1 100644
--- a/onnx/test/parser_test.py
+++ b/onnx/test/parser_test.py
@@ -66,3 +66,7 @@ def test_parse_model_error(self) -> None:
            }
            '''
         self.assertRaises(onnx.parser.ParseError, lambda: onnx.parser.parse_model(input))
+
+
+if __name__ == '__main__':
+    unittest.main()


Code review comments:
diff_hunk: @@ -77,3 +68,24 @@
     int(OptionalProto.MAP): 'map_value',
     int(OptionalProto.OPTIONAL): 'optional_value'
 }
+
+
+def to_np_type(tensor_type: str) -> Any:
path: onnx/mapping.py
text: @user1:
Are the `str` annotations correct? (Here and several places below.)

diff_hunk: @@ -2,52 +2,42 @@
 
 from onnx import TensorProto, SequenceProto, OptionalProto
 import numpy as np  # type: ignore
+from typing import Any
+import warnings
 
-# This map is used for converting TensorProto values into Numpy arrays
-TENSOR_TYPE_TO_NP_TYPE = {
-    int(TensorProto.FLOAT): np.dtype('float32'),
-    int(TensorProto.UINT8): np.dtype('uint8'),
-    int(TensorProto.INT8): np.dtype('int8'),
-    int(TensorProto.UINT16): np.dtype('uint16'),
-    int(TensorProto.INT16): np.dtype('int16'),
-    int(TensorProto.INT32): np.dtype('int32'),
-    int(TensorProto.INT64): np.dtype('int64'),
-    int(TensorProto.BOOL): np.dtype('bool'),
-    int(TensorProto.FLOAT16): np.dtype('float16'),
-    int(TensorProto.BFLOAT16): np.dtype('float32'),  # Native numpy does not support bfloat16 so now use float32 for bf16 values
-    int(TensorProto.DOUBLE): np.dtype('float64'),
-    int(TensorProto.COMPLEX64): np.dtype('complex64'),
-    int(TensorProto.COMPLEX128): np.dtype('complex128'),
-    int(TensorProto.UINT32): np.dtype('uint32'),
-    int(TensorProto.UINT64): np.dtype('uint64'),
-    int(TensorProto.STRING): np.dtype('object')
+# tesnor_type: (numpy type, storage type, string name)
path: onnx/mapping.py
text: @user1:
nit: spelling "tensor"

diff_hunk: @@ -77,3 +68,24 @@
     int(OptionalProto.MAP): 'map_value',
     int(OptionalProto.OPTIONAL): 'optional_value'
 }
+
+
+def to_np_type(tensor_type: str) -> Any:
+    return TENSOR_TYPE_MAP[tensor_type][0]
+
+
+def to_sotrage_tensor_type(tensor_type: str) -> int:
+    return TENSOR_TYPE_MAP[tensor_type][1]
+
+
+def to_string(tensor_type: str) -> str:
path: onnx/mapping.py
text: @user1:
nit: May be useful to include the "from" type also in the name. May be `tensor_dtype_to_string` or something like that?

@user1:
Suggesting `tensor_dtype` as `tensor_type` is somewhat overloaded already. We mean the element-type of tensors, right?

@author:
Updated. I used `tensor_dtype_to_XXX` for naming. Thanks!

diff_hunk: @@ -77,3 +68,24 @@
     int(OptionalProto.MAP): 'map_value',
     int(OptionalProto.OPTIONAL): 'optional_value'
 }
+
+
+def to_np_type(tensor_type: str) -> Any:
+    return TENSOR_TYPE_MAP[tensor_type][0]
+
+
+def to_sotrage_tensor_type(tensor_type: str) -> int:
path: onnx/mapping.py
text: @user1:
nit: spelling "storage"