Key op fixes for failing tests

Author: gs-olive

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

@@ -2,6 +2,7 @@
 from typing import Any, List, Optional, Sequence, Union, cast
 
 import numpy as np
+import tensorrt as trt
 import torch
 from torch.fx.node import Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
@@ -19,8 +20,6 @@
 from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
 from torch_tensorrt.fx.utils import get_dynamic_dims
 
-import tensorrt as trt
-
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
 
@@ -101,9 +100,15 @@ def layer_norm(
             "of the TensorRT region!"
         )
 
-    gamma = weight.detach().cpu().float().numpy()
+    gamma = (
+        weight.detach().cpu().float().numpy()
+        if isinstance(weight, torch.Tensor)
+        else weight
+    )
     gamma_field = trt.PluginField("gamma", gamma, trt.PluginFieldType.FLOAT32)
-    beta = bias.detach().cpu().float().numpy()
+    beta = (
+        bias.detach().cpu().float().numpy() if isinstance(bias, torch.Tensor) else bias
+    )
     beta_field = trt.PluginField("beta", beta, trt.PluginFieldType.FLOAT32)
     eps_field = trt.PluginField(
         "eps", np.array(eps, dtype=np.float32), trt.PluginFieldType.FLOAT32

py/torch_tensorrt/fx/converters/acc_ops_converters.py

@@ -3,30 +3,27 @@
 import math
 import operator
 import warnings
-from typing import cast, Dict, Optional, Sequence, Tuple, Union
+from typing import Dict, Optional, Sequence, Tuple, Union, cast
 
 import numpy as np
 
 # @manual=//deeplearning/trt/python:py_tensorrt
 import tensorrt as trt
 import torch
-
-from ..converter_registry import tensorrt_converter
-
-from ..tracer.acc_tracer import acc_ops
-from ..types import *  # noqa: F403
 from torch.fx.immutable_collections import immutable_list
 from torch.fx.node import Argument, Target
-
-from ..utils import get_dynamic_dims, unified_dtype_converter, Frameworks
-
-from .converter_utils import *  # noqa: F403
+from torch_tensorrt.fx.converters.impl import activation, convolution
 from torch_tensorrt.fx.passes.lower_basic_pass import (
     trt_transposed_linear,
     trt_transposed_matmul,
 )
 from torch_tensorrt.fx.tracer.acc_tracer.acc_ops import contiguous
-from torch_tensorrt.fx.converters.impl import activation, convolution
+
+from ..converter_registry import tensorrt_converter
+from ..tracer.acc_tracer import acc_ops
+from ..types import *  # noqa: F403
+from ..utils import Frameworks, get_dynamic_dims, unified_dtype_converter
+from .converter_utils import *  # noqa: F403
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -2714,8 +2711,14 @@ def acc_ops_linear(
         "dim for linear and it can't be the last dim."
     )
 
-    if isinstance(kwargs["weight"], torch.Tensor):
-        weight = get_trt_tensor(network, kwargs["weight"].t(), f"{name}_weight")
+    if isinstance(kwargs["weight"], (torch.Tensor, np.ndarray)):
+        weight = get_trt_tensor(
+            network,
+            kwargs["weight"].t()
+            if isinstance(kwargs["weight"], torch.Tensor)
+            else kwargs["weight"].T,
+            f"{name}_weight",
+        )
         if target not in (acc_ops.linear, torch.ops.aten.linear):
             weight_op = trt.MatrixOperation.TRANSPOSE
         else:

py/torch_tensorrt/fx/converters/converter_utils.py

@@ -1,8 +1,8 @@
 import operator
 import warnings
+from enum import Enum, auto
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
 
-from enum import Enum, auto
 import numpy as np
 
 # @manual=//deeplearning/trt/python:py_tensorrt
@@ -20,7 +20,7 @@
     TRTPluginFieldCollection,
     TRTTensor,
 )
-from ..utils import unified_dtype_converter, Frameworks
+from ..utils import Frameworks, unified_dtype_converter
 
 
 class SourceIR(Enum):
@@ -271,7 +271,7 @@ def create_constant(
     """
     constant = network.add_constant(
         (1,) if isinstance(value, (int, float)) else value.shape,
-        to_numpy(value, dtype),
+        to_numpy(value, dtype).copy(),
     )
     constant.name = name
     return constant.get_output(0)
@@ -311,7 +311,7 @@ def get_trt_tensor(
     elif isinstance(input_val, np.ndarray) and (
         input_val.dtype == np.bool_ or input_val.dtype == np.int64
     ):
-        input_val = input_val.to(np.int32)
+        input_val = input_val.astype(np.int32)
 
     if isinstance(input_val, (torch.Tensor, np.ndarray, int, float)):
         return create_constant(network, input_val, name, dtype)

py/torch_tensorrt/fx/converters/impl/convolution.py

@@ -1,27 +1,23 @@
-import numpy as np
 from typing import Any, Optional, Sequence, Union
 
+import numpy as np
+
 # @manual=//deeplearning/trt/python:py_tensorrt
 import tensorrt as trt
 import torch
 from torch.fx.node import Target
-
+from torch_tensorrt.fx.converters import acc_ops_converters
 from torch_tensorrt.fx.converters.converter_utils import (
     SourceIR,
     extend_attr_to_tuple,
     get_dyn_range,
+    get_trt_tensor,
+    has_dynamic_shape,
     mark_as_int8_layer,
     set_layer_name,
-    has_dynamic_shape,
     to_numpy,
-    get_trt_tensor,
-)
-from torch_tensorrt.fx.converters import acc_ops_converters
-
-from torch_tensorrt.fx.types import (
-    TRTNetwork,
-    TRTTensor,
 )
+from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
 
 
 def convNd(
@@ -54,7 +50,7 @@ def convNd(
         )
 
     # Process bias terms
-    if isinstance(bias, torch.Tensor):
+    if isinstance(bias, (torch.Tensor, np.ndarray)):
         # Transform the bias constant into a Numpy array
         bias = to_numpy(bias)
 
@@ -79,7 +75,7 @@ def convNd(
                 network, target, tuple(), kwargs, name + "_unsqueeze_weight"
             )
 
-    elif isinstance(weight, torch.Tensor):
+    elif isinstance(weight, (torch.Tensor, np.ndarray)):
         # Transform the weight constant into a Numpy array
         weight = to_numpy(weight)