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/converter_utils.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
@@ -414,3 +413,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,
+) -> 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.
+    """
+    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):
+        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/impl/conv.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/deconv.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/normalization/ops.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/select.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/shape.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