Key op fixes for failing tests

Author: gs-olive

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

@@ -112,9 +112,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/dynamo/conversion/impl/permutation.py

@@ -1,4 +1,5 @@
-from typing import Optional, Sequence, cast
+from typing import Optional, Sequence, Union, cast
+import numpy as np
 
 
 from torch.fx.node import Target
@@ -8,6 +9,7 @@
 from torch_tensorrt.fx.converters.converter_utils import (
     set_layer_name,
     get_positive_dim,
+    get_trt_tensor,
 )
 
 
@@ -16,9 +18,15 @@ def permute(
     target: Target,
     source_ir: Optional[SourceIR],
     name: str,
-    input: TRTTensor,
+    input: Union[TRTTensor, np.ndarray],
     permutation: Sequence[int],
 ) -> TRTTensor:
+    if isinstance(input, np.ndarray):
+        tensor_to_freeze = np.transpose(input, permutation)
+        # TODO: Fix naming for constant tensors
+        frozen_trt_tensor = get_trt_tensor(network, tensor_to_freeze, name)
+        return frozen_trt_tensor
+
     if not isinstance(input, TRTTensor):
         raise RuntimeError(
             f"permute received input {input} that is not a TensorRT ITensor"

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

@@ -1,5 +1,6 @@
 from typing import Optional, cast
 import math
+import numpy as np
 
 from torch.fx.node import Target
 
@@ -25,12 +26,6 @@ def slice_op(
     stop: int,
     step: int,
 ) -> TRTTensor:
-    if not isinstance(input, TRTTensor):
-        raise RuntimeError(
-            f"slice_tensor received input {input} that is not part "
-            "of the TensorRT region!"
-        )
-
     ranks = len(input.shape) + (1 if network.has_implicit_batch_dimension else 0)
     dim = get_positive_dim(cast(int, dim), ranks)
     dynamic_shape = has_dynamic_shape(input.shape)
@@ -49,6 +44,22 @@ def slice_op(
     if stop_int == 2**63 - 1:
         stop_int = input.shape[dim]
     step_int = cast(int, step)
+
+    if isinstance(input, np.ndarray):
+        tensor_to_freeze = np.take(
+            input, np.arange(start_int, stop_int, step_int), axis=dim
+        )
+        # TODO: Fix naming for constant tensors
+        frozen_trt_tensor = get_trt_tensor(network, tensor_to_freeze, name)
+        return frozen_trt_tensor
+
+    if not isinstance(input, TRTTensor):
+        raise RuntimeError(
+            f"slice_tensor received input {input} that is not part "
+            "of the TensorRT region!"
+        )
+
+    # TRT Input Formatting
     start = [0] * len(input.shape)
     start[dim] = start_int
     stride = [1] * len(start)

py/torch_tensorrt/fx/converters/acc_ops_converters.py

@@ -2719,8 +2719,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

@@ -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)

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

@@ -55,7 +55,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)
 
@@ -80,7 +80,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)