added the flattening function to util

Author: cehongwang

py/torch_tensorrt/dynamo/utils.py

@@ -3,7 +3,7 @@
 import logging
 from dataclasses import fields, replace
 from enum import Enum
-from typing import Any, Callable, Dict, Optional, Sequence, Union
+from typing import Any, Callable, Dict, Generator, List, Optional, Sequence, Union
 
 import numpy as np
 import tensorrt as trt
@@ -413,32 +413,29 @@ def check_output(
     return True
 
 
-def unified_dtype_converter(
-    dtype: Union[TRTDataType, torch.dtype, np.dtype], to: Frameworks
-) -> Union[np.dtype, torch.dtype, TRTDataType]:
+def flatten_dict_value(d: dict[Any, Any]) -> List[Any]:
     """
-    Convert TensorRT, Numpy, or Torch data types to any other of those data types.
+    Flatten the values of a dictionary to a single list.
 
     Args:
-        dtype (TRTDataType, torch.dtype, np.dtype): A TensorRT, Numpy, or Torch data type.
-        to (Frameworks): The framework to convert the data type to.
+    d (dict): The dictionary to flatten.
 
     Returns:
-        The equivalent data type in the requested framework.
+    list: A list of all values flattened.
     """
-    assert to in Frameworks, f"Expected valid Framework for translation, got {to}"
-    trt_major_version = int(trt.__version__.split(".")[0])
-    if dtype in (np.int8, torch.int8, trt.int8):
-        return DataTypeEquivalence[trt.int8][to]
-    elif trt_major_version >= 7 and dtype in (np.bool_, torch.bool, trt.bool):
-        return DataTypeEquivalence[trt.bool][to]
-    elif dtype in (np.int32, torch.int32, trt.int32):
-        return DataTypeEquivalence[trt.int32][to]
-    elif dtype in (np.int64, torch.int64, trt.int64):
-        return DataTypeEquivalence[trt.int64][to]
-    elif dtype in (np.float16, torch.float16, trt.float16):
-        return DataTypeEquivalence[trt.float16][to]
-    elif dtype in (np.float32, torch.float32, trt.float32):
-        return DataTypeEquivalence[trt.float32][to]
-    else:
-        raise TypeError("%s is not a supported dtype" % dtype)
+
+    def flatten(value: Any) -> Generator[Any, Any, Any]:
+        if isinstance(value, dict):
+            for v in value.values():
+                yield from flatten(v)
+        elif isinstance(value, list):
+            for item in value:
+                yield from flatten(item)
+        else:
+            yield value
+
+    flat_list: List[Any] = []
+    for v in d.values():
+        flat_list.extend(flatten(v))
+
+    return flat_list