fix: Upgrade to_numpy to allow boolean constants

Author: gs-olive

py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py

@@ -345,7 +345,7 @@ def call_function(self, target: str, args: Any, kwargs: Any) -> Any:
 
     def get_attr(self, target: str, args: Any, kwargs: Any) -> np.ndarray:
         with _disable_current_modes():
-            from torch_tensorrt.fx.converters import to_numpy
+            from torch_tensorrt.dynamo.conversion.converter_utils import to_numpy
 
             frozen_attr = self.fetch_attr(target)
 

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -17,7 +17,6 @@
 from torch_tensorrt.fx.converters.converter_utils import (
     Frameworks,
     get_axes_for_reduce_op,
-    to_numpy,
     unified_dtype_converter,
 )
 from torch_tensorrt.fx.types import TRTDataType, TRTTensor
@@ -270,9 +269,10 @@ def create_constant(
     Returns:
         A TensorRT ITensor that represents the given value.
     """
+    numpy_value = to_numpy(value, dtype)
     constant = ctx.net.add_constant(
         (1,) if isinstance(value, (int, float, bool)) else value.shape,
-        to_numpy(value, dtype).copy(),
+        numpy_value.copy() if isinstance(numpy_value, np.ndarray) else numpy_value,
     )
     constant.name = name
     return constant.get_output(0)
@@ -351,7 +351,10 @@ def enforce_tensor_types(
             and all((dtype in (TRTTensor, np.ndarray, torch.Tensor)) for dtype in val)
         )
         for val in type_dictionary.values()
-    ), "Invalid value(s) specified in type enforcement"
+    ), (
+        "Invalid value(s) specified in type enforcement."
+        "Note that torch.Tensor cannot be present as a type without np.ndarray."
+    )
 
     def wrapper(func: DynamoConverterImplSignature) -> DynamoConverterImplSignature:
         @functools.wraps(func)
@@ -384,20 +387,25 @@ def convert_with_type_enforcement(
                         f"Detected argument at index {index} had type {type(candidate)} "
                         f"which is not one of the approved types {approved_dtypes}"
                     )
-                # Numpy arrays are preferred in general - if approved, promote to Numpy first
-                elif np.ndarray in approved_dtypes and not isinstance(
-                    candidate, TRTTensor
-                ):
-                    new_value = to_numpy(candidate)
-                # As a fallback, freeze tensors to IConstantLayers if they cannot be handled as Numpy arrays
-                elif TRTTensor in approved_dtypes:
-                    _LOGGER.debug(
-                        f"Freezing tensor {name}_constant_{index} to TRT IConstantLayer"
-                    )
-                    new_value = get_trt_tensor(
-                        ctx, candidate, name + f"_constant_{index}"
-                    )
-                else:
+
+                # Type-promotion preference order depends on tuple order
+                for dtype in approved_dtypes:
+                    # Currently, we do not cast to Torch tensor, due to issues with such casts
+                    # in FakeTensor contexts
+                    if dtype == np.ndarray and not isinstance(candidate, TRTTensor):
+                        new_value = to_numpy(candidate)
+                        break
+                    # As a fallback, freeze tensors to IConstantLayers if they cannot be handled as Numpy arrays
+                    elif dtype == TRTTensor:
+                        _LOGGER.debug(
+                            f"Freezing tensor {name}_constant_{index} to TRT IConstantLayer"
+                        )
+                        new_value = get_trt_tensor(
+                            ctx, candidate, name + f"_constant_{index}"
+                        )
+                        break
+
+                if new_value is None:
                     raise AssertionError(
                         f"Argument at index {index} was not able to be converted to one of "
                         f"the following types: {approved_dtypes}"
@@ -414,3 +422,50 @@ def convert_with_type_enforcement(
         return convert_with_type_enforcement
 
     return wrapper
+
+
+def to_numpy(
+    value: Optional[Union[torch.Tensor, np.ndarray, int, float, bool]],
+    dtype: Optional[Union[torch.dtype, np.dtype, TRTDataType]] = None,
+) -> Optional[np.ndarray]:
+    """
+    Convert a PyTorch Tensor, Numpy array, or scalar to a Numpy Array. If the tensor is
+    quantized it will be dequantized first.
+    Args:
+        value (Optional[Union[torch.Tensor, np.ndarray, int, float, bool]]):
+            A PyTorch tensor, Numpy array, int, float, or bool
+        dtype (Optional[Union[torch.dtype, np.dtype, TRTDataType]]):
+            If a dtype is given, we will convert the type of the given `value` to this dtype.
+    Returns:
+        A Numpy array or None, if the input was None.
+    """
+    output = None
+
+    if value is None or isinstance(value, np.ndarray):
+        output = value
+
+    elif isinstance(value, torch.Tensor):
+        if value.is_quantized:
+            value = value.dequantize()
+
+        output = value.cpu().detach().contiguous().numpy()
+
+    elif isinstance(value, int):
+        output = np.array([value], dtype=np.int32)
+
+    elif isinstance(value, float):
+        output = np.array([value], dtype=np.float32)
+
+    elif isinstance(value, bool):
+        output = np.array([value], dtype=np.bool_)
+
+    if isinstance(output, np.ndarray) or output is None:
+        return (
+            output
+            if (dtype is None or output is None)
+            else output.astype(unified_dtype_converter(dtype, Frameworks.NUMPY))
+        )
+    else:
+        raise AssertionError(
+            f"to_numpy can only be called on None, bool, int, float, np.ndarray, or torch.Tensor, got: {value}"
+        )

py/torch_tensorrt/dynamo/conversion/converter_utils.py

@@ -9,16 +9,16 @@
 from torch_tensorrt.dynamo.conversion import impl
 from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
 from torch_tensorrt.dynamo.conversion.converter_utils import (
+    SourceIR,
     extend_attr_to_tuple,
     get_trt_tensor,
+    to_numpy,
 )
 from torch_tensorrt.fx.converters.converter_utils import (
-    SourceIR,
     get_dyn_range,
     has_dynamic_shape,
     mark_as_int8_layer,
     set_layer_name,
-    to_numpy,
 )
 from torch_tensorrt.fx.types import TRTTensor
 

py/torch_tensorrt/dynamo/conversion/impl/conv.py

@@ -11,14 +11,14 @@
 from torch_tensorrt.dynamo.conversion.converter_utils import (
     extend_attr_to_tuple,
     get_trt_tensor,
+    to_numpy,
 )
 from torch_tensorrt.fx.converters.converter_utils import (
     SourceIR,
     get_dyn_range,
     has_dynamic_shape,
     mark_as_int8_layer,
     set_layer_name,
-    to_numpy,
 )
 from torch_tensorrt.fx.types import TRTTensor
 

py/torch_tensorrt/dynamo/conversion/impl/deconv.py

@@ -7,6 +7,7 @@
 from torch.fx.node import Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
 from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
+from torch_tensorrt.dynamo.conversion.converter_utils import to_numpy
 from torch_tensorrt.dynamo.conversion.impl.elementwise.base import (
     convert_binary_elementwise,
 )
@@ -16,7 +17,6 @@
     get_trt_plugin,
     has_dynamic_shape,
     set_layer_name,
-    to_numpy,
 )
 from torch_tensorrt.fx.types import TRTTensor
 from torch_tensorrt.fx.utils import get_dynamic_dims

py/torch_tensorrt/dynamo/conversion/impl/normalization/ops.py

@@ -3,16 +3,17 @@
 import tensorrt as trt
 from torch.fx.node import Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
+from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
 from torch_tensorrt.dynamo.conversion.converter_utils import extend_attr_to_tuple
 from torch_tensorrt.fx.converters.converter_utils import (
     has_dynamic_shape,
     set_layer_name,
 )
-from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
+from torch_tensorrt.fx.types import TRTTensor
 
 
 def avg_poolNd(
-    network: TRTNetwork,
+    ctx: ConversionContext,
     target: Union[Target, str],
     source_ir: Optional[SourceIR],
     name: str,
@@ -45,7 +46,7 @@ def avg_poolNd(
     padding = extend_attr_to_tuple(padding, dim)
 
     # add average pooling layer
-    pool_layer = network.add_pooling_nd(
+    pool_layer = ctx.net.add_pooling_nd(
         input=input,
         type=trt.PoolingType.AVERAGE,
         window_size=kernel_size,
@@ -60,7 +61,7 @@ def avg_poolNd(
 
 
 def max_poolNd(
-    network: TRTNetwork,
+    ctx: ConversionContext,
     target: Union[Target, str],
     source_ir: Optional[SourceIR],
     name: str,
@@ -92,7 +93,7 @@ def max_poolNd(
     padding = extend_attr_to_tuple(padding, dim)
 
     # add max pooling layer
-    pool_layer = network.add_pooling_nd(
+    pool_layer = ctx.net.add_pooling_nd(
         input=input,
         type=trt.PoolingType.MAX,
         window_size=kernel_size,

py/torch_tensorrt/dynamo/conversion/impl/pool.py

@@ -4,11 +4,11 @@
 from torch.fx.node import Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
 from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
+from torch_tensorrt.dynamo.conversion.converter_utils import to_numpy
 from torch_tensorrt.dynamo.conversion.impl.shape import get_shape_with_dynamic_shape
 from torch_tensorrt.fx.converters.converter_utils import (
     get_positive_dim,
     has_dynamic_shape,
-    to_numpy,
 )
 from torch_tensorrt.fx.types import Shape, TRTTensor
 

py/torch_tensorrt/dynamo/conversion/impl/select.py

@@ -8,10 +8,11 @@
 from torch.fx.node import Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
 from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
+from torch_tensorrt.dynamo.conversion.converter_utils import to_numpy
 from torch_tensorrt.dynamo.conversion.impl.elementwise.base import (
     convert_binary_elementwise,
 )
-from torch_tensorrt.fx.converters.converter_utils import set_layer_name, to_numpy
+from torch_tensorrt.fx.converters.converter_utils import set_layer_name
 from torch_tensorrt.fx.types import TRTTensor