chore: import upstream snapshot with attribution

tensorflow/compiler/tf2tensorrt/convert/ops/binary_ops.cc

@@ -0,0 +1,235 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+const BinaryOperationMapType* BinaryOperationMap() {
+  static const auto* map = new BinaryOperationMapType({
+    {"Add", nvinfer1::ElementWiseOperation::kSUM},
+        {"AddV2", nvinfer1::ElementWiseOperation::kSUM},
+        {"Mul", nvinfer1::ElementWiseOperation::kPROD},
+        {"Sub", nvinfer1::ElementWiseOperation::kSUB},
+        {"Div", nvinfer1::ElementWiseOperation::kDIV},
+        {"FloorDiv", nvinfer1::ElementWiseOperation::kFLOOR_DIV},
+        {"RealDiv", nvinfer1::ElementWiseOperation::kDIV},
+        {"Minimum", nvinfer1::ElementWiseOperation::kMIN},
+        {"Maximum", nvinfer1::ElementWiseOperation::kMAX},
+        {"Pow", nvinfer1::ElementWiseOperation::kPOW},
+#if IS_TRT_VERSION_GE(8, 2, 0, 0)
+        {"Greater", nvinfer1::ElementWiseOperation::kGREATER},
+        {"Less", nvinfer1::ElementWiseOperation::kLESS},
+        {"Equal", nvinfer1::ElementWiseOperation::kEQUAL},
+        // Operators are implemented as NOT Less and NOT Greater, respectively.
+        {"GreaterEqual", nvinfer1::ElementWiseOperation::kLESS},
+        {"LessEqual", nvinfer1::ElementWiseOperation::kGREATER},
+#endif
+  });
+  return map;
+}
+
+const BinaryOperationMapType* BinaryBooleanOperationMap() {
+  static const auto* map = new BinaryOperationMapType({
+      {"LogicalOr", nvinfer1::ElementWiseOperation::kOR},
+      {"LogicalAnd", nvinfer1::ElementWiseOperation::kAND},
+  });
+  return map;
+}
+
+namespace {
+class ConvertBinaryImpl {
+ protected:
+  ConvertBinaryImpl(const BinaryOperationMapType* pOperMap)
+      : pOperMap_(pOperMap) {}
+
+  Status ValidateImpl(
+      const OpConverterParams& params,
+      const std::vector<string>& implicit_batch_not_supported_ops = {},
+      bool both_tensors = false) {
+    const auto& node_def = params.node_def;
+    const auto& op = node_def.op();
+    const auto op_pair = pOperMap_->find(op);
+    if (op_pair == pOperMap_->end()) {
+      return errors::Unimplemented("Binary op: ", op, " not supported");
+    }
+
+    // Constant folding should have been done by TensorFlow.
+    const auto& inputs = params.inputs;
+    if (inputs.at(0).is_weights() && inputs.at(1).is_weights()) {
+      return errors::Unimplemented(
+          "Constant folding is falled back to TensorFlow, binary op '", op,
+          "' received both input as constant");
+    }
+
+    if ((convertToBool_ = find_name(op, implicit_batch_not_supported_ops))) {
+      if (params.use_implicit_batch) {
+        return errors::Unimplemented(
+            convert_not_supported_implicit(op, node_def.name(), "Binary"));
+      }
+    }
+
+    if (both_tensors) {
+      if (inputs.at(0).is_weights() || inputs.at(1).is_weights()) {
+        return errors::InvalidArgument("Both inputs  of '", op,
+                                       "' are expected to be tensors");
+      }
+      // No need to convert the output of "LogicalOr" and "LogicalAnd"
+      convertToBool_ = false;
+    }
+
+    nvinfer1::Dims broadcasted_dims[2];
+    TF_RETURN_IF_ERROR(GetTrtBroadcastShape(
+        inputs.at(0), inputs.at(1), true, params.use_implicit_batch,
+        broadcasted_dims, broadcasted_dims + 1));
+
+    for (int i = 0; i < tensor_.size(); i++) {
+      // This will also convert constants to tensors.
+      TF_RETURN_IF_ERROR(PrepareTensorForShape(
+          params.converter, inputs.at(i), broadcasted_dims[i],
+          params.validation_only, &tensor_[i], node_def, i));
+    }
+    operation_ = op_pair->second;
+    return OkStatus();
+  }
+
+  Status ConvertImpl(const OpConverterParams& params,
+                     const std::vector<string>& revert_bool_ops = {}) {
+    const auto& node_def = params.node_def;
+    // Add ElementWise layer.
+    auto* network = params.converter->network();
+    nvinfer1::ILayer* layer = network->addElementWise(
+        *tensor_[0]->trt_tensor(), *tensor_[1]->trt_tensor(), operation_);
+    TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name());
+
+    if (params.use_explicit_precision) {
+      layer->setPrecision(nvinfer1::DataType::kFLOAT);
+    }
+
+    params.converter->SetLayerName(layer, node_def);
+    const auto& output = layer->getOutput(0);
+    if (convertToBool_) {
+      output->setType(nvinfer1::DataType::kBOOL);
+      if (find_name(node_def.op(), revert_bool_ops)) {
+        nvinfer1::IUnaryLayer* unary_layer =
+            network->addUnary(*output, nvinfer1::UnaryOperation::kNOT);
+        TFTRT_RETURN_ERROR_IF_NULLPTR(unary_layer, node_def.name());
+        params.outputs->push_back(
+            TRT_TensorOrWeights(unary_layer->getOutput(0)));
+        return OkStatus();
+      }
+    }
+
+    params.outputs->push_back(TRT_TensorOrWeights(output));
+    return OkStatus();
+  }
+
+  static constexpr std::array<InputArgSpec, 2> InputSpec() {
+    return std::array<InputArgSpec, 2>{
+        InputArgSpec::Create("x", TrtInputArg::kBoth),
+        InputArgSpec::Create("y", TrtInputArg::kBoth)};
+  }
+
+ private:
+  const BinaryOperationMapType* pOperMap_;
+  std::array<ITensorProxyPtr, 2> tensor_{nullptr, nullptr};
+  nvinfer1::ElementWiseOperation operation_;
+  bool convertToBool_;
+};
+
+class ConvertBinary : public OpConverterBase<ConvertBinary>,
+                      protected ConvertBinaryImpl {
+ public:
+  explicit ConvertBinary(const OpConverterParams* params)
+      : OpConverterBase<ConvertBinary>(
+            params,
+            {DataType::DT_FLOAT, DataType::DT_HALF, DataType::DT_INT32}),
+        ConvertBinaryImpl(BinaryOperationMap()) {}
+
+  static constexpr std::array<InputArgSpec, 2> InputSpec() {
+    return ConvertBinaryImpl::InputSpec();
+  }
+
+  Status Validate() {
+    const std::vector<string> implicit_batch_not_supported_ops {
+#if IS_TRT_VERSION_GE(8, 2, 0, 0)
+      "Greater", "Less", "Equal", "GreaterEqual", "LessEqual"
+#endif
+    };
+    return ValidateImpl(*params_, implicit_batch_not_supported_ops);
+  }
+  Status Convert() {
+    const std::vector<string> implemented_with_reverted_ops {
+#if IS_TRT_VERSION_GE(8, 2, 0, 0)
+      "GreaterEqual", "LessEqual"
+#endif
+    };
+    return ConvertImpl(*params_, implemented_with_reverted_ops);
+  }
+};
+
+class ConvertBooleanBinary : public OpConverterBase<ConvertBooleanBinary>,
+                             public ConvertBinaryImpl {
+ public:
+  explicit ConvertBooleanBinary(const OpConverterParams* params)
+      : OpConverterBase<ConvertBooleanBinary>(params, {DataType::DT_BOOL}),
+        ConvertBinaryImpl(BinaryBooleanOperationMap()) {}
+
+  static constexpr std::array<InputArgSpec, 2> InputSpec() {
+    return ConvertBinaryImpl::InputSpec();
+  }
+
+  static constexpr const char* NodeDefDataTypeAttributeName() {
+    /*
+    node {
+      name: "..."
+      op: "LogicalOr"
+      input: "..."
+      input: "..."
+      attr {
+        key: "_output_shapes"
+        ...
+      }
+    }
+    */
+    return "";
+  }
+  Status Validate() {
+#if IS_TRT_VERSION_GE(8, 2, 0, 0)
+    return ValidateImpl(*params_, {"LogicalOr", "LogicalAnd"}, true);
+#else
+    return errors::Unimplemented("Boolean op: ", params_->node_def.op(),
+                                 " is not supported in TRT version < 8.2");
+#endif
+  }
+  Status Convert() { return ConvertImpl(*params_); }
+};
+}  // namespace
+
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertBinary>(),
+                                  GetOperationNames(*BinaryOperationMap()));
+REGISTER_DEFAULT_TRT_OP_CONVERTER(
+    MakeConverterFunction<ConvertBooleanBinary>(),
+    GetOperationNames(*BinaryBooleanOperationMap()));
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/data_format_vec_permute.cc

@@ -0,0 +1,176 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "third_party/tensorrt/NvInfer.h"
+#include "third_party/tensorrt/NvInferRuntimeCommon.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+int get_spatial_dim_count(string format) {
+  // Spatial dimensions are the dimensions besides NC, and here we assume NC
+  // always appear in the format string.
+  return format.size() - 2;
+}
+
+class ConvertDataFormatVecPermute
+    : public OpConverterBase<ConvertDataFormatVecPermute> {
+ public:
+  ConvertDataFormatVecPermute(const OpConverterParams* params)
+      : OpConverterBase<ConvertDataFormatVecPermute>(params,
+                                                     {DataType::DT_INT32}) {}
+
+  struct DataFormatVecPermuteAttributes {
+    string dst_format;
+    string src_format;
+    int x_dim_count;
+  };
+
+  static constexpr std::array<InputArgSpec, 1> InputSpec() {
+    return {InputArgSpec::Create("x", TrtInputArg::kBoth)};
+  }
+
+  Status Validate() {
+    TF_RETURN_IF_ERROR(NotSupportedInImplicitBatch());
+    const auto& inputs = params_->inputs;
+    const auto& nodeName = params_->node_def.name();
+
+    x_input_ = inputs.at(0);
+
+    // Check input rank.
+    const auto x_dims = x_input_.GetTrtDims();
+    int input_rank = x_dims.nbDims;
+    if (input_rank != 1 && input_rank != 2) {
+      return errors::InvalidArgument(
+          "Input must be a vector or matrix, but got rank ", input_rank,
+          ", at ", nodeName);
+    }
+
+    // Verify and consume node attributes.
+    StatusOr<string> dst_format = GetAttrValue<string>("dst_format");
+    StatusOr<string> src_format = GetAttrValue<string>("src_format");
+    TRT_ENSURE_OK(dst_format);
+    TRT_ENSURE_OK(src_format);
+
+    // Check input dims.
+    const int full_dim_count = src_format->size();
+    const int spatial_dim_count = get_spatial_dim_count(*src_format);
+    if (input_rank == 1) {
+      if (x_dims.d[0] != spatial_dim_count && x_dims.d[0] != full_dim_count) {
+        return errors::InvalidArgument(
+            "1D input must be of size ", spatial_dim_count, " or ",
+            full_dim_count, ", but got size ", x_dims.d[0], ", at ", nodeName);
+      }
+    } else if (input_rank == 2) {
+      if (x_dims.d[0] != spatial_dim_count && x_dims.d[0] != full_dim_count) {
+        return errors::InvalidArgument(
+            "First dimension of 2D input must be of size ", spatial_dim_count,
+            " or ", full_dim_count, ", but got shape (", x_dims.d[0], ", ",
+            x_dims.d[1], "), at ", nodeName);
+      }
+      if (x_dims.d[1] != 2) {
+        return errors::InvalidArgument(
+            "Second dimension of 2D input must be of size 2, but got shape (",
+            x_dims.d[0], ", ", x_dims.d[1], "), at ", nodeName);
+      }
+    }
+
+    // Set custom attributes.
+    attrs_.x_dim_count = x_dims.d[0];
+    attrs_.dst_format = *dst_format;
+    attrs_.src_format = *src_format;
+
+    return OkStatus();
+  }
+
+  Status Convert() {
+    // Copy format strings in case they need to be modified.
+    string dst_format = attrs_.dst_format;
+    string src_format = attrs_.src_format;
+    const int& spatial_dim_count = get_spatial_dim_count(src_format);
+
+    // If the input is a vector of size spatial_dim_count, treat the elements
+    // as spatial dimensions.
+    if (attrs_.x_dim_count == spatial_dim_count) {
+      auto keep_only_spatial_dimensions =
+          [spatial_dim_count](string* format_str) -> void {
+        auto new_end = std::remove_if(format_str->begin(), format_str->end(),
+                                      [spatial_dim_count](const char dim) {
+                                        return dim == 'N' || dim == 'C';
+                                      });
+        format_str->erase(new_end, format_str->end());
+      };
+      keep_only_spatial_dimensions(&src_format);
+      keep_only_spatial_dimensions(&dst_format);
+    }
+
+    // Create indices for the gather layer and make weights out of them.
+    std::vector<int32> dst_indices(attrs_.x_dim_count);
+    for (int i = 0; i < attrs_.x_dim_count; ++i) {
+      for (int j = 0; j < attrs_.x_dim_count; ++j) {
+        if (src_format[i] == dst_format[j]) {
+          dst_indices[j] = i;
+          break;
+        }
+      }
+    }
+    nvinfer1::Dims indices_dims = {1, {attrs_.x_dim_count}};
+    StatusOr<TRT_ShapedWeights> indices_weights =
+        params_->weight_store->GetTempWeights(nvinfer1::DataType::kINT32,
+                                              indices_dims);
+    TRT_ENSURE_OK(indices_weights);
+    int32* indices_ptr = indices_weights->GetPointer<int32>();
+    std::copy(dst_indices.data(), dst_indices.data() + attrs_.x_dim_count,
+              indices_ptr);
+    ITensorProxyPtr x_tensor =
+        x_input_.is_weights() ? params_->converter->CreateConstantLayer(
+                                    x_input_.weights(), x_input_.GetTrtDims())
+                              : x_input_.tensor();
+    ITensorProxyPtr indices_tensor =
+        params_->converter->CreateConstantLayer(*indices_weights, indices_dims);
+
+    // Gather layer with 1D indices on axis 0, conserves shape.
+    nvinfer1::IGatherLayer* layer = params_->converter->network()->addGather(
+        *x_tensor->trt_tensor(), *indices_tensor->trt_tensor(), 0);
+    TRT_ENSURE(layer);
+    params_->converter->SetLayerName(layer, params_->node_def);
+
+    ITensorProxyPtr output_tensor = layer->getOutput(0);
+
+    params_->outputs->push_back(TRT_TensorOrWeights(output_tensor));
+    return OkStatus();
+  }
+
+ private:
+  TRT_TensorOrWeights x_input_;
+  DataFormatVecPermuteAttributes attrs_{};
+};
+REGISTER_DEFAULT_TRT_OP_CONVERTER(
+    MakeConverterFunction<ConvertDataFormatVecPermute>(),
+    {"DataFormatVecPermute"});
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/einsum.cc

@@ -0,0 +1,908 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#include <iterator>
+#include <limits>
+#include <memory>
+
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/util/einsum_op_util.h"
+#include "third_party/tensorrt/NvInfer.h"
+
+#if IS_TRT_VERSION_GE(7, 1, 3, 0)
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+namespace {
+
+#if !IS_TRT_VERSION_GE(8, 2, 0, 0)
+
+// Finds the indices of elements in [begin, end) in array
+// [array_begin, array_end), and appends the indices to permute. This is used to
+// construct the permutation sequence for the operand with input labels
+// [array_begin, array_end) to the desired permuted labels [begin, end).
+template <typename T>
+Status FindIndicesoOfAllValuesInSrc(absl::Span<const T> values,
+                                    absl::Span<const T> src,
+                                    std::vector<int>* indices) {
+  if (src.size() < values.size()) {
+    return errors::Internal(
+        "Span 'src' cannot contain all elements of 'values'");
+  }
+  for (auto i = 0; i < values.size(); i++) {
+    auto iter = absl::c_find(src, values[i]);
+    if (iter == src.end()) {
+      return errors::Internal("Label ", values[i], " not found");
+    }
+    int idx = std::distance(src.begin(), iter);
+    indices->push_back(idx);
+  }
+  return OkStatus();
+}
+
+// Layout of the einsum dimensions: Batch, Free and Contraction indices.
+// Example: adbc,adce -> adbe. The first tensor has layout BFC, the second BCF.
+enum class EinsumLayout { BFC, BCF, MIX };
+
+using DimType = EinsumDimensionType;
+constexpr auto kBatch = DimType::kBatch;
+constexpr auto kFree = DimType::kFree;
+constexpr auto kContract = DimType::kContract;
+
+// Describes an operand: input shape, number of batch, free and contract
+// dimensions, and the permutation that is needed to bring it to a matmul
+// compatible form.
+class EinsumDescriptor {
+ private:
+  // Checks whether input_labels[offset:offset+m] matches labels from other.
+  static bool OrderMatches(const Labels& input_labels, int offset, int m,
+                           EinsumDimensionType dim_type,
+                           const std::unique_ptr<EinsumDescriptor>& other) {
+    if (other == nullptr) {
+      return true;
+    }
+    int offset_other = 0;
+    if (dim_type == kFree) {
+      offset = other->offset_f;
+    } else if (dim_type == kContract) {
+      offset = other->offset_c;
+    }
+    return std::equal(input_labels.begin() + offset,
+                      input_labels.begin() + offset + m,
+                      other->permuted_labels.begin() + offset_other);
+  }
+
+  using label_t_iterator = std::vector<EinsumDimensionType>::const_iterator;
+  static int32_t CountLabels(label_t_iterator begin, label_t_iterator end,
+                             EinsumDimensionType val) {
+    return static_cast<int32_t>(std::count_if(
+        begin, end, [val](EinsumDimensionType t) { return t == val; }));
+  }
+
+  // Appends indices to the "permute" vector where types matches value.
+  void AppendMatchingIndicesToPermute(
+      const std::vector<EinsumDimensionType>& types, EinsumDimensionType val) {
+    for (int i = 0; i < types.size(); i++) {
+      if (types[i] == val) {
+        permute.push_back(i);
+      }
+    }
+  }
+
+  Status DetermineLayout(const Labels& input_labels,
+                         const std::vector<EinsumDimensionType>& types,
+                         const std::unique_ptr<EinsumDescriptor>& other) {
+    // Check if the current layout is BFC or BCF. In that case we could avoid
+    // transpose.
+    layout = EinsumLayout::MIX;
+    if (CountLabels(types.begin(), types.begin() + b, kBatch) == b &&
+        OrderMatches(input_labels, 0, b, kBatch, other)) {
+      // Batch dims are the leading dims. They have the same order as other.
+      if (CountLabels(types.begin() + b, types.begin() + b + f, kFree) == f) {
+        // All the free dims are placed consecutively after the batch dims.
+        // Their order is arbitrary. The final transpose will ensure that the
+        // output has correct order. We still have to check that the contract
+        // indices have correct order.
+        if (OrderMatches(input_labels, b + f, c, kContract, other)) {
+          layout = EinsumLayout::BFC;
+        }
+      } else if (CountLabels(types.begin() + b, types.begin() + b + c,
+                             kContract) == c) {
+        // All the contract dims are placed consecutively after the batch
+        // dims. Check whether the contract dims have the same order as the
+        // contract dims in other.
+        if (OrderMatches(input_labels, b, c, kContract, other)) {
+          layout = EinsumLayout::BCF;
+        }
+      }
+    }
+    return OkStatus();
+  }
+
+  Status CalculateMixedLayoutPermutation(
+      const EinsumLayout preferred_layout, const Labels& input_labels,
+      const std::vector<EinsumDimensionType>& types,
+      const std::unique_ptr<EinsumDescriptor>& other) {
+    // Input label types are mixed. Calculate a permutation that maps them
+    // to the preferred layout (BCF or BFC).
+    layout = preferred_layout;
+    if (other == nullptr) {
+      AppendMatchingIndicesToPermute(types, kBatch);
+    } else {
+      TF_RETURN_IF_ERROR(
+          FindIndicesoOfAllValuesInSrc(/*values=*/
+                                       absl::MakeConstSpan(
+                                           other->permuted_labels.begin(),
+                                           other->b),
+                                       /*src=*/
+                                       absl::MakeConstSpan(input_labels.begin(),
+                                                           input_labels.size()),
+                                       /*indices=*/&permute));
+    }
+    if (layout == EinsumLayout::BFC) {
+      AppendMatchingIndicesToPermute(types, kFree);
+      if (other == nullptr) {
+        AppendMatchingIndicesToPermute(types, kContract);
+      } else {
+        TF_RETURN_IF_ERROR(FindIndicesoOfAllValuesInSrc(
+            /*values=*/absl::MakeConstSpan(
+                other->permuted_labels.begin() + other->offset_c, other->c),
+            /*src=*/
+            absl::MakeConstSpan(input_labels.begin(), input_labels.size()),
+            /*indices=*/&permute));
+      }
+      return OkStatus();
+    }
+    if (other == nullptr) {
+      AppendMatchingIndicesToPermute(types, kContract);
+    } else {
+      TF_RETURN_IF_ERROR(FindIndicesoOfAllValuesInSrc(
+          /*values=*/absl::MakeConstSpan(
+              other->permuted_labels.begin() + other->offset_c, other->c),
+          /*src=*/absl::MakeConstSpan(input_labels.begin(), input_labels.end()),
+          /*indices=*/&permute));
+    }
+    AppendMatchingIndicesToPermute(types, kFree);
+    return OkStatus();
+  }
+
+  Status Initialize(const TRT_TensorOrWeights& operand, Labels input_labels,
+                    std::vector<EinsumDimensionType>& label_types,
+                    EinsumLayout preferred_layout,
+                    const std::unique_ptr<EinsumDescriptor>& other = nullptr) {
+    if (preferred_layout == EinsumLayout::MIX) {
+      return errors::Internal("Preferred einsum layout cannot be MIX");
+    }
+    // Map label indices to label types.
+    std::vector<EinsumDimensionType> types;  // Input label types.
+    std::transform(input_labels.begin(), input_labels.end(),
+                   std::back_inserter(types),
+                   [&label_types](int i) { return label_types.at(i); });
+
+    b = CountLabels(types.begin(), types.end(), kBatch);
+    f = CountLabels(types.begin(), types.end(), kFree);
+    c = CountLabels(types.begin(), types.end(), kContract);
+
+    if (c == 0 || f == 0) {
+      VLOG(2) << "Einsum equation needs to have at least one free and one "
+                 "contract dimension";
+      return errors::Unimplemented("No conversion for einsum equation.");
+    }
+
+    TF_RETURN_IF_ERROR(DetermineLayout(input_labels, types, other));
+    if (layout == EinsumLayout::MIX) {
+      TF_RETURN_IF_ERROR(CalculateMixedLayoutPermutation(
+          preferred_layout, input_labels, types, other));
+    }
+
+    if (layout == EinsumLayout::BFC) {
+      offset_f = b;
+      offset_c = f + b;
+    } else {
+      offset_f = b + c;
+      offset_c = b;
+    }
+
+    dims = operand.GetTrtDims();
+    for (int i = 0; i < b; i++) {
+      // Set unknown batch dims to zero. These dims will be used in reshape op,
+      // where zero is a special value for retaining the original dim size.
+      if (dims.d[i] == -1) {
+        dims.d[i] = 0;
+      }
+    }
+    permuted_labels = input_labels;
+    if (!permute.empty()) {
+      // Apply the permutation on the dimension array.
+      nvinfer1::Dims orig_dims = dims;
+      for (int i = 0; i < permute.size(); i++) {
+        dims.d[i] = orig_dims.d[permute[i]];
+        permuted_labels[i] = input_labels[permute[i]];
+      }
+    }
+    size_tensors.resize(dims.nbDims, nullptr);
+    return OkStatus();
+  }
+
+ public:
+  EinsumDescriptor() : b(0), f(0), c(0) {}
+
+  // Deduces the number of batch, free, contract dimensions from the input
+  // labels, decides what layout to use, and determines permutation indices for
+  // that layout.
+  static StatusOr<std::unique_ptr<EinsumDescriptor>> Create(
+      const TRT_TensorOrWeights& operand, Labels input_labels,
+      std::vector<EinsumDimensionType>& label_types,
+      EinsumLayout preferred_layout,
+      const std::unique_ptr<EinsumDescriptor>& other = nullptr) {
+    auto desc = std::make_unique<EinsumDescriptor>();
+    TF_RETURN_IF_ERROR(desc->Initialize(operand, input_labels, label_types,
+                                        preferred_layout, other));
+    VLOG(2) << desc->DebugString();
+    return desc;
+  }
+
+  int NumBatchDims() const { return b; }
+  int NumContractDims() const { return c; }
+  int NumFreeDims() const { return f; }
+  int ContractDimOffset() const { return offset_c; }
+  const Labels& PermutedLabels() const { return permuted_labels; }
+
+  std::string DebugString() const {
+    return absl::StrCat("Descriptor with ",
+                        (layout == EinsumLayout::BFC ? "BFC" : "BCF"),
+                        " layout, b=", b, ", f=", f, ", c=", c);
+  }
+
+  // Returns whether the free and contract dimension have static shape.
+  bool HasStaticShape() const {
+    return !std::any_of(dims.d + b, dims.d + dims.nbDims,
+                        [](int k) { return k == -1; });
+  }
+
+  nvinfer1::Permutation GetPermutation() const {
+    nvinfer1::Permutation p;
+    std::copy(permute.begin(), permute.end(), p.order);
+    return p;
+  }
+
+  std::vector<int> PermuteVector() const { return permute; }
+
+  // Sets the "size_tensors" vector to be filled with scalar constant tensors
+  // representing the shape of the operand.
+  Status SetDynamicSize(TRTNetworkBuilder* builder,
+                        const TRT_TensorOrWeights& operand) {
+    TRT_ENSURE(operand.GetTrtDims().nbDims == dims.nbDims);
+    if (operand.is_weights()) {
+      // Generate constants for each dimension of the constant weight tensor's
+      // shape.
+      for (int i = 0; i < operand.GetTrtDims().nbDims; i++) {
+        StatusOr<nvinfer1::IConstantLayer*> size_tensor =
+            builder->Constant<int32_t>(dims.d[i], 1);
+        TRT_ENSURE_PTR_OK(size_tensor);
+        size_tensors[i] = (*size_tensor)->getOutput(0);
+      }
+      return OkStatus();
+    }
+
+    // If the operand is a dynamic tensor, compute the shape value dynamically.
+    StatusOr<nvinfer1::IShapeLayer*> shape_layer =
+        builder->Shape(operand.tensor()->trt_tensor());
+    TRT_ENSURE_PTR_OK(shape_layer);
+    nvinfer1::ITensor* shape = (*shape_layer)->getOutput(0);
+    for (int i = 0; i < operand.GetTrtDims().nbDims; i++) {
+      int idx = permute.empty() ? i : permute.at(i);
+      StatusOr<nvinfer1::ISliceLayer*> slice_layer =
+          builder->Slice(shape, {1, {idx}}, {1, {1}}, {1, {1}});
+      TRT_ENSURE_PTR_OK(slice_layer);
+      size_tensors[i] = (*slice_layer)->getOutput(0);
+    }
+    return OkStatus();
+  }
+
+  EinsumLayout layout;
+  int b;  // number of batch dims
+  int f;  // number of free dims
+  int c;  // number of conraction dims
+  int offset_f;
+  int offset_c;
+  nvinfer1::Dims dims;
+  std::vector<int> permute;
+  std::vector<ITensorProxyPtr> size_tensors;
+  Labels permuted_labels;
+};
+
+// Reshapes operand so that the free dimensions are combined into a single dim,
+// and the contract dimensions are combined into another single dim.
+Status GetEinsumNewDynamicShape(TRTNetworkBuilder* builder,
+                                const EinsumDescriptor& desc,
+                                ITensorProxyPtr* new_shape) {
+  std::vector<nvinfer1::ITensor*> size;
+  size.reserve(desc.b + 2);
+  absl::c_transform(absl::MakeSpan(desc.size_tensors).subspan(0, desc.b + 2),
+                    std::back_inserter(size),
+                    [](const ITensorProxyPtr x) { return x->trt_tensor(); });
+
+  int idx_f = desc.layout == EinsumLayout::BFC ? desc.b : desc.b + 1;
+  int idx_c = desc.layout == EinsumLayout::BFC ? desc.b + 1 : desc.b;
+
+  std::vector<nvinfer1::ITensor*> size_tensors;
+  size_tensors.reserve(desc.size_tensors.size());
+  absl::c_transform(desc.size_tensors, std::back_inserter(size_tensors),
+                    [](const ITensorProxyPtr x) -> nvinfer1::ITensor* {
+                      return x->trt_tensor();
+                    });
+
+  StatusOr<nvinfer1::ILayer*> shape_vol = builder->CumulativeProd(
+      absl::MakeSpan(size_tensors).subspan(desc.offset_f, desc.f));
+  TRT_ENSURE_PTR_OK(shape_vol);
+  size[idx_f] = (*shape_vol)->getOutput(0);
+
+  shape_vol = builder->CumulativeProd(
+      absl::MakeSpan(size_tensors).subspan(desc.offset_c, desc.c));
+  TRT_ENSURE_PTR_OK(shape_vol);
+  size[idx_c] = (*shape_vol)->getOutput(0);
+  StatusOr<nvinfer1::IConcatenationLayer*> layer =
+      builder->Concat(size, /*axis=*/0);
+  TRT_ENSURE_PTR_OK(layer);
+  *new_shape = (*layer)->getOutput(0);
+  return OkStatus();
+}
+
+// Reshapes operand so that the free dimensions are combined into a single dim,
+// and the contract dimensions are combined into another single dim.
+Status GetEinsumNewStaticShape(const EinsumDescriptor& desc,
+                               nvinfer1::Dims* new_dims) {
+  // Copy the batch dims and append two additional dimensions.
+  DimsAdapter adap(
+      absl::MakeSpan(static_cast<const int32_t*>(desc.dims.d), desc.b));
+  adap.Append(1).Append(1);
+
+  // Combine free dims and contract dims.
+  int idx_f = desc.layout == EinsumLayout::BFC ? desc.b : desc.b + 1;
+  int idx_c = desc.layout == EinsumLayout::BFC ? desc.b + 1 : desc.b;
+
+  // Find the volume of the free dimensions.
+  int64_t vol_f =
+      DimsAdapter(
+          absl::MakeSpan(
+              static_cast<const int32_t*>(desc.dims.d) + desc.offset_f, desc.f))
+          .Volume();
+
+  // Find the volume of the contracted dimensions.
+  int64_t vol_c =
+      DimsAdapter(
+          absl::MakeSpan(
+              static_cast<const int32_t*>(desc.dims.d) + desc.offset_c, desc.c))
+          .Volume();
+
+  adap.dim(idx_f) = vol_f;
+  adap.dim(idx_c) = vol_c;
+  *new_dims = adap.AsTrtDims();
+  return OkStatus();
+}
+
+StatusOr<TRT_TensorOrWeights> ConditionEinsumWeights(
+    TRTNetworkBuilder* builder, const TRT_TensorOrWeights& operand,
+    const EinsumDescriptor& desc, const bool need_transpose) {
+  TRT_ENSURE(operand.is_weights());
+  if (!need_transpose) {
+    // If we don't need to transpose, then the operand remains as a weights
+    // constant. In this case we also don't need a reshape.
+    TRT_ShapedWeights weights(operand.weights());
+    nvinfer1::Dims new_dims;
+    TF_RETURN_IF_ERROR(GetEinsumNewStaticShape(desc, &new_dims));
+    TF_RETURN_IF_ERROR(weights.SetShape(new_dims));
+    return TRT_TensorOrWeights(weights);
+  }
+
+  // Let TensorRT handle constant folding where possible.
+  StatusOr<nvinfer1::IConstantLayer*> tensor = builder->WeightsToConstant(
+      operand.weights().GetTrtWeights(), operand.GetTrtDims());
+  TRT_ENSURE_PTR_OK(tensor);
+  return TRT_TensorOrWeights((*tensor)->getOutput(0));
+}
+
+// Builds a TRT shuffle operation for the given operand. Replaces operand with a
+// pointer to the shuffle output.
+Status ConditionEinsumTensor(TRTNetworkBuilder* builder,
+                             std::unique_ptr<TRT_TensorOrWeights>* operand,
+                             const EinsumDescriptor& desc,
+                             const bool need_transpose,
+                             const bool need_reshape) {
+  StatusOr<ShuffleBuilder> shuffle =
+      ShuffleBuilder::Create(builder, (*operand)->tensor()->trt_tensor());
+  TRT_ENSURE_OK(shuffle);
+
+  // Set new shape.
+  if (need_reshape) {
+    if (desc.HasStaticShape()) {
+      nvinfer1::Dims new_dims;
+      TF_RETURN_IF_ERROR(GetEinsumNewStaticShape(desc, &new_dims));
+      shuffle->SetReshape(new_dims);
+    } else {
+      ITensorProxyPtr new_shape;
+      TF_RETURN_IF_ERROR(GetEinsumNewDynamicShape(&*builder, desc, &new_shape));
+      shuffle->SetReshape(new_shape->trt_tensor());
+    }
+  }
+
+  if (need_transpose) {
+    shuffle->SetFirstTranspose(desc.GetPermutation());
+  }
+
+  StatusOr<nvinfer1::ITensor*> shuffle_out = shuffle->Output();
+  TRT_ENSURE_PTR_OK(shuffle_out);
+  *operand = std::make_unique<TRT_TensorOrWeights>(*shuffle_out);
+  return OkStatus();
+}
+
+// Handles einsum operand conditioning for both constant and non-constant
+// inputs. This is supported using the ShuffleEinsumWeights and
+// ShuffleEinsumTensor routines.
+Status ConditionEinsumOperand(TRTNetworkBuilder* builder,
+                              std::unique_ptr<TRT_TensorOrWeights>* operand,
+                              const EinsumDescriptor& desc) {
+  bool need_reshape = (desc.f != 1 || desc.c != 1);
+  bool need_transpose = !desc.permute.empty();
+
+  VLOG(2) << "Condition operand. Need reshape: " << need_reshape
+          << ". Need transpose: " << need_transpose;
+
+  if ((*operand)->is_weights()) {
+    StatusOr<TRT_TensorOrWeights> result =
+        ConditionEinsumWeights(builder, **operand, desc, need_transpose);
+    TRT_ENSURE_OK(result);
+    *operand = std::make_unique<TRT_TensorOrWeights>(std::move(result).value());
+  }
+
+  // If we didn't convert the operand to a tensor, we can return here.
+  if ((*operand)->is_weights()) {
+    return OkStatus();
+  }
+
+  TF_RETURN_IF_ERROR(ConditionEinsumTensor(builder, operand, desc,
+                                           need_transpose, need_reshape));
+
+  return OkStatus();
+}
+
+// Combines output dims/labels by copying batch and free dims/labels from input
+// A, and concatenating free values from input B.
+template <typename InputIterator, typename OutputIterator>
+void AssembleOutput(InputIterator begin_a, InputIterator begin_b,
+                    const EinsumDescriptor& desc_a,
+                    const EinsumDescriptor& desc_b, OutputIterator out) {
+  std::copy(begin_a, begin_a + desc_a.b, out);
+  begin_a += desc_a.offset_f;
+  std::copy(begin_a, begin_a + desc_a.f, out + desc_a.b);
+  begin_b += desc_b.offset_f;
+  std::copy(begin_b, begin_b + desc_b.f, out + desc_a.b + desc_a.f);
+}
+
+// Restores free dimensions and sets final index order. Consider C = A * B,
+// batched MatMul op, where A.shape = [B, x, k] and B.shape = [B, k, y]. Then
+// C.shape = [B, x, y]. Here B can denote multiple batch indices while x, y, k
+// are single indices. The original inputs to Einsum can have multiple free
+// indices. These were combined into a singe free dimension x and y, for example
+// x = f_a1 * f_a2 * f_a3, y = f_b1 * f_b2. This routine creates a shuffle layer
+// to expand x into and y the original free dims, e.g. C is reshaped to
+// [B, f_a1, f_a2, f_a3, f_b1, f_b2]. Finally, a permutation is applied to
+// transform the shape to the shape of the original Einsum output.
+Status ShuffleEinsumOutput(const OpConverterParams* params,
+                           EinsumDescriptor desc_a, EinsumDescriptor desc_b,
+                           const std::vector<int>& permutation,
+                           ITensorProxyPtr* output) {
+  if (permutation.empty() && (desc_a.f == 1 && desc_b.f == 1)) {
+    return OkStatus();
+  }
+
+  nvinfer1::IShuffleLayer* layer =
+      params->converter->network()->addShuffle(*(*output)->trt_tensor());
+  TFTRT_RETURN_ERROR_IF_NULLPTR(layer, params->node_def.name());
+  params->converter->SetLayerName(layer, params->node_def, "shuffle",
+                                  /*sub_op_instance=*/2);
+
+  int output_rank = desc_a.b + desc_a.f + desc_b.f;
+  if (desc_a.f != 1 || desc_b.f != 1) {
+    if (desc_a.HasStaticShape() && desc_b.HasStaticShape()) {
+      nvinfer1::Dims dims_out = {output_rank, {}};
+      AssembleOutput(desc_a.dims.d, desc_b.dims.d, desc_a, desc_b, dims_out.d);
+      layer->setReshapeDimensions(dims_out);
+    } else {
+      std::vector<ITensorProxyPtr> size_tensors(output_rank);
+      AssembleOutput(desc_a.size_tensors.begin(), desc_b.size_tensors.begin(),
+                     desc_a, desc_b, size_tensors.begin());
+      ITensorProxyPtr new_shape;
+      auto builder = TRTNetworkBuilder::Create(params->converter->network(),
+                                               params->weight_store);
+      TRT_ENSURE_OK(builder);
+      std::vector<nvinfer1::ITensor*> size_itensors;
+      absl::c_transform(size_tensors, std::back_inserter(size_itensors),
+                        [](auto x) { return x->trt_tensor(); });
+      StatusOr<nvinfer1::IConcatenationLayer*> concat =
+          builder->Concat(size_itensors, /*axis=*/0);
+      TRT_ENSURE_PTR_OK(concat);
+      new_shape = (*concat)->getOutput(0);
+      layer->setInput(1, *new_shape->trt_tensor());
+    }
+  }
+
+  if (!permutation.empty()) {
+    nvinfer1::Permutation p;
+    std::copy(permutation.begin(), permutation.end(), p.order);
+    layer->setSecondTranspose(p);
+  }
+  *output = layer->getOutput(0);
+  return OkStatus();
+}
+
+// Updates "final_transpose" according to the given descriptors and output
+// labels.
+StatusOr<std::vector<int>> GetOutputTranspose(
+    const EinsumDescriptor& descriptor_a, const EinsumDescriptor& descriptor_b,
+    Labels output_labels) {
+  // Get final transpose.
+  std::vector<int> final_transpose;
+  final_transpose.reserve(descriptor_a.b + descriptor_a.f + descriptor_b.f);
+  Labels matmul_output_labels(descriptor_a.b + descriptor_a.f + descriptor_b.f);
+  AssembleOutput(descriptor_a.permuted_labels.begin(),
+                 descriptor_b.permuted_labels.begin(), descriptor_a,
+                 descriptor_b, matmul_output_labels.begin());
+  TF_RETURN_IF_ERROR(
+      FindIndicesoOfAllValuesInSrc(/*values=*/
+                                   absl::MakeConstSpan(output_labels.begin(),
+                                                       output_labels.end()),
+                                   /*src=*/
+                                   absl::MakeConstSpan(
+                                       matmul_output_labels.begin(),
+                                       matmul_output_labels.end()),
+                                   /*indices=*/&final_transpose));
+  // Clear identity transpose.
+  bool identity_transpose = true;
+  for (int i = 0; i < final_transpose.size() && identity_transpose; i++) {
+    identity_transpose &= final_transpose.at(i) == i;
+  }
+  if (identity_transpose) {
+    final_transpose.clear();
+  }
+  return final_transpose;
+}
+
+// Prepares EinsumDescriptors after parsing the equation and determines the
+// final transpose.
+Status ParseEquation(const std::string& equation,
+                     std::unique_ptr<TRT_TensorOrWeights>* input_a,
+                     std::unique_ptr<TRT_TensorOrWeights>* input_b,
+                     std::unique_ptr<EinsumDescriptor>* descriptor_a,
+                     std::unique_ptr<EinsumDescriptor>* descriptor_b,
+                     std::vector<int>* final_transpose) {
+  VLOG(2) << "Einsum equation " << equation;
+  OperandLabels input_labels;
+  Labels output_labels;
+  std::vector<EinsumDimensionType> label_types;
+  OperandLabelCounts input_label_counts;
+  LabelCounts output_label_counts;
+  absl::InlinedVector<bool, 2> input_has_ellipsis;
+  bool output_has_ellipsis;
+  TF_RETURN_IF_ERROR(
+      ParseEinsumEquation(equation, &input_labels, &output_labels, &label_types,
+                          &input_label_counts, &output_label_counts,
+                          &input_has_ellipsis, &output_has_ellipsis));
+
+  if (input_has_ellipsis[0] || input_has_ellipsis[1] || output_has_ellipsis) {
+    // TODO(tfeher): Handle ellipsis like EinsumHelper::ProcessDimensions.
+    // Note: ProcessDimensions would introduce kBroadcasting labels, which we
+    // need to replace with kBatch before we call InitDescriptor.
+    VLOG(2) << "Ellipsis not yet supported";
+    return errors::Unimplemented("No conversion for einsum equation.");
+  }
+
+  if (absl::c_any_of(label_types, [](auto l) {
+        return l == EinsumDimensionType::kReduce ||
+               l == EinsumDimensionType::kBroadcasting;
+      })) {
+    VLOG(2) << "Einsum reductions not implemented";
+    return errors::Unimplemented("No conversion for einsum equation.");
+  }
+
+  auto no_duplicated_labels = [](const LabelCounts& label_counts) {
+    return absl::c_any_of(label_counts, [](int i) { return i > 1; });
+  };
+  if (no_duplicated_labels(input_label_counts[0]) ||
+      no_duplicated_labels(input_label_counts[1]) ||
+      no_duplicated_labels(output_label_counts)) {
+    VLOG(2) << "Einsum invalid label count";
+    return errors::Unimplemented("No conversion for einsum equation.");
+  }
+
+  if ((*input_a)->is_weights() && (*input_b)->is_tensor()) {
+    // We prefer to use FC layer, needs A as tensor and B as weight.
+    std::swap(*input_a, *input_b);
+    std::swap(input_labels[0], input_labels[1]);
+    std::swap(input_label_counts[0], input_label_counts[1]);
+  }
+
+  auto desc = EinsumDescriptor::Create(**input_a, input_labels[0], label_types,
+                                       EinsumLayout::BFC);
+  TF_RETURN_IF_ERROR(desc.status());
+  *descriptor_a = std::move(desc).value();
+
+  desc = EinsumDescriptor::Create(**input_b, input_labels[1], label_types,
+                                  EinsumLayout::BCF, *descriptor_a);
+  TF_RETURN_IF_ERROR(desc.status());
+  *descriptor_b = std::move(desc).value();
+
+  auto out_transpose =
+      GetOutputTranspose(**descriptor_a, **descriptor_b, output_labels);
+
+  TRT_ENSURE_OK(out_transpose)
+  *final_transpose = std::move(out_transpose).value();
+  return OkStatus();
+}
+
+class ConvertEinsum : public OpConverterBase<ConvertEinsum> {
+ public:
+  explicit ConvertEinsum(const OpConverterParams* params)
+      : OpConverterBase<ConvertEinsum>(params) {}
+
+  static constexpr std::array<InputArgSpec, 2> InputSpec() {
+    return {InputArgSpec::Create("input_a", TrtInputArg::kBoth),
+            InputArgSpec::Create("input_b", TrtInputArg::kBoth)};
+  }
+
+  Status Validate() {
+    TF_RETURN_IF_ERROR(NotSupportedInImplicitBatch());
+    const auto& inputs = params_->inputs;
+    input_a = std::make_unique<TRT_TensorOrWeights>(inputs.at(0));
+    input_b = std::make_unique<TRT_TensorOrWeights>(inputs.at(1));
+
+    StatusOr<std::string> eq = GetAttrValue<std::string>("equation");
+    TRT_ENSURE_OK(eq);
+    TF_RETURN_IF_ERROR(ParseEquation(*eq, &input_a, &input_b, &descriptor_a,
+                                     &descriptor_b, &final_transpose));
+    return OkStatus();
+  }
+
+  Status Convert() {
+    auto builder = TRTNetworkBuilder::Create(params_->converter->network(),
+                                             params_->weight_store);
+    TRT_ENSURE_OK(builder);
+    TRT_ENSURE(input_a && input_b);
+    TRT_ENSURE(descriptor_a && descriptor_b);
+
+    // Populate the size_tensor vector in the descriptor.
+    TF_RETURN_IF_ERROR(descriptor_a->SetDynamicSize(&*builder, *input_a));
+    TF_RETURN_IF_ERROR(descriptor_b->SetDynamicSize(&*builder, *input_b));
+
+    // Condition the operands for lowering to matmul.
+    TF_RETURN_IF_ERROR(
+        ConditionEinsumOperand(&*builder, &input_a, *descriptor_a));
+    TF_RETURN_IF_ERROR(
+        ConditionEinsumOperand(&*builder, &input_b, *descriptor_b));
+
+    // Build the matmul implementation.
+    StatusOr<ITensorProxyPtr> result = ConvertMatMulImpl(
+        params_, *input_a, *input_b, descriptor_a->layout == EinsumLayout::BCF,
+        descriptor_b->layout == EinsumLayout::BFC);
+    TF_RETURN_IF_ERROR(result.status());
+    ITensorProxyPtr output = result.value();
+
+    // Reshape and permute the output.
+    TF_RETURN_IF_ERROR(ShuffleEinsumOutput(
+        params_, *descriptor_a, *descriptor_b, final_transpose, &output));
+    this->AddOutput(output);
+    return OkStatus();
+  }
+
+ private:
+  std::unique_ptr<TRT_TensorOrWeights> input_a{nullptr};
+  std::unique_ptr<TRT_TensorOrWeights> input_b{nullptr};
+  std::vector<int> final_transpose;
+  std::unique_ptr<EinsumDescriptor> descriptor_a{nullptr};
+  std::unique_ptr<EinsumDescriptor> descriptor_b{nullptr};
+};
+#else
+
+// Helper class to reindex equations to contain only lowercase characters. We
+// simply define a mapping from the old character set to a new set.
+// - The input is assumed to be a valid TF equation.
+// - The input is TRT compatible, therefore it has max 8 dims. (Thus we have
+//   enough lowercase English characters to represent the equation.)
+// How do we reindex/map equations:
+// - Only uppercase letters are changed, if possible we just lowercase them.
+// - If the equation contains both upper and lowercase variant of a letter, say
+//   X and x, then we map X to the first unused lowercase letter.
+class ReIndexer {
+ public:
+  // Initializes the index map with existing lowercase labels.
+  ReIndexer(std::string eq) {
+    for (char c : eq) {
+      if (absl::ascii_islower(c)) {
+        idx_map_[c] = c;
+      }
+    }
+  }
+  // Finds new character for uppercase character c.
+  char operator()(char c) {
+    if (!absl::ascii_isupper(c)) return c;
+    if (idx_map_.count(c) > 0) return idx_map_[c];
+    char new_idx = absl::ascii_tolower(c);
+
+    // If lower(c) is not used in the equation, use it to replace c.
+    if (idx_map_.count(new_idx) == 0) {
+      idx_map_[c] = new_idx;
+      idx_map_[new_idx] = new_idx;  // mark that new_idx is taken
+      return new_idx;
+    }
+
+    // Otherwise, find the first available lower case to replace c.
+    for (char k = 'a'; k <= 'z'; k++) {
+      if (idx_map_.count(k) == 0) {
+        new_idx = k;
+        idx_map_[c] = new_idx;
+        idx_map_[new_idx] = new_idx;  // mark that new_idx is taken
+        break;
+      }
+    }
+    return new_idx;
+  }
+
+ private:
+  // Each key is an index used in the original or in the reindexed equation.
+  // The values are the corresponding new lowercase indices.
+  std::map<char, char> idx_map_;
+};
+
+class ConvertEinsum : public OpConverterBase<ConvertEinsum> {
+ public:
+  explicit ConvertEinsum(const OpConverterParams* params)
+      : OpConverterBase<ConvertEinsum>(params) {}
+
+  Status ValidateInputs() {
+    TRT_ENSURE(params_->inputs.size() <= 2);
+    return OkStatus();
+  }
+  static constexpr bool HasFixNumberOfInputs() { return false; }
+
+  static constexpr std::array<InputArgSpec, 2> InputSpec() {
+    return {InputArgSpec::Create("input_a", TrtInputArg::kBoth),
+            InputArgSpec::Create("input_b", TrtInputArg::kBoth)};
+  }
+
+  std::string MakeLowerCase(const std::string& eq) {
+    std::string res = eq;
+    ReIndexer reindexer(eq);
+    std::transform(eq.begin(), eq.end(), res.begin(), reindexer);
+    return res;
+  }
+
+  // Checks if the equation is supported by TRT.
+  Status ValidateEinsumEquation(const std::string& eq) {
+    const auto& inputs = params_->inputs;
+    OperandLabels input_labels;
+    Labels output_labels;
+    std::vector<EinsumDimensionType> label_types;
+    OperandLabelCounts input_label_counts;
+    LabelCounts output_label_counts;
+    absl::InlinedVector<bool, 2> input_has_ellipsis;
+    bool output_has_ellipsis;
+    VLOG(2) << "Parsing equation " << eq;
+    TF_RETURN_IF_ERROR(ParseEinsumEquation(
+        eq, &input_labels, &output_labels, &label_types, &input_label_counts,
+        &output_label_counts, &input_has_ellipsis, &output_has_ellipsis));
+
+    Status unimplemented =
+        errors::Unimplemented("No conversion for einsum equation.");
+    if (input_has_ellipsis[0] || (inputs.size() > 1 && input_has_ellipsis[1]) ||
+        output_has_ellipsis) {
+      VLOG(2) << "Ellipsis not yet supported";
+      return unimplemented;
+    }
+    for (int i = 0; i < input_label_counts.size(); i++) {
+      for (int k = 0; k < input_label_counts[i].size(); k++) {
+        if (input_label_counts[i][k] > 1) {
+          VLOG(2) << "Diagonal operation or reduction not yet supported";
+          return unimplemented;
+        }
+      }
+    }
+    bool has_out_idx =
+        std::reduce(output_label_counts.begin(), output_label_counts.end(),
+                    false, std::logical_or<int>());
+    if (!has_out_idx) {
+      VLOG(2) << "Scalar output not allowed in dynamic shape mode";
+      return unimplemented;
+    }
+    // Check for outer product
+    if (input_label_counts.size() == 2 && output_label_counts.size() == 2 &&
+        output_label_counts[0] == 1 && output_label_counts[1] == 1) {
+      VLOG(2) << "Outer product not supported";
+      return unimplemented;
+    }
+    return OkStatus();
+  }
+
+  Status Validate() {
+    VLOG(2) << "Running validation using the new einsum "
+               "converter";
+    TF_RETURN_IF_ERROR(NotSupportedInImplicitBatch());
+    StatusOr<std::string> eq = GetAttrValue<std::string>("equation");
+    TRT_ENSURE_OK(eq);
+
+    TF_RETURN_IF_ERROR(ValidateEinsumEquation(*eq));
+
+    // While TF has case sensitive equations, TensorRT expects lowercase eq (as
+    // of version 8.4). See
+    // https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#einsum-layer
+    equation = MakeLowerCase(*eq);
+
+    return OkStatus();
+  }
+
+  Status Convert() {
+    auto builder = TRTNetworkBuilder::Create(params_->converter->network(),
+                                             params_->weight_store);
+    TRT_ENSURE_OK(builder);
+
+    std::vector<nvinfer1::ITensor*> trt_input;
+    for (const TRT_TensorOrWeights& input_arg : params_->inputs) {
+      ITensorProxyPtr ptr = nullptr;
+      if (input_arg.is_tensor()) {
+        ptr = input_arg.tensor();
+      } else {
+        StatusOr<nvinfer1::IConstantLayer*> const_layer =
+            builder->WeightsToConstant(input_arg.weights().GetTrtWeights(),
+                                       input_arg.GetTrtDims());
+        TRT_ENSURE_PTR_OK(const_layer);
+        ptr = (*const_layer)->getOutput(0);
+      }
+      trt_input.push_back(ptr->trt_tensor());
+    }
+    nvinfer1::IEinsumLayer* layer = params_->converter->network()->addEinsum(
+        trt_input.data(), trt_input.size(), equation.c_str());
+    TRT_ENSURE(layer);
+
+    ITensorProxyPtr output = layer->getOutput(0);
+    this->AddOutput(output);
+    return OkStatus();
+  }
+
+ private:
+  std::string equation;
+};
+
+#endif
+
+}  // namespace
+
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertEinsum>(),
+                                  "Einsum");
+#endif
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/fill_ops.cc

@@ -0,0 +1,313 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+#if IS_TRT_VERSION_GE(8, 2, 0, 0)
+
+template <typename Impl>
+class ConvertFillBase : public OpConverterBase<Impl> {
+ public:
+  explicit ConvertFillBase(const OpConverterParams* params)
+      : OpConverterBase<Impl>(params, {DataType::DT_FLOAT, DataType::DT_HALF,
+                                       DataType::DT_INT32}) {}
+};
+
+class ConvertFill : public ConvertFillBase<ConvertFill> {
+ public:
+  explicit ConvertFill(const OpConverterParams* params)
+      : ConvertFillBase<ConvertFill>(params) {}
+
+  static constexpr std::array<InputArgSpec, 2> InputSpec() {
+    return std::array<InputArgSpec, 2>{
+        InputArgSpec::Create("dims", TrtInputArg::kBoth),
+        InputArgSpec::Create("value", TrtInputArg::kBoth)};
+  }
+
+  Status Validate() {
+    const auto& params = *this->params_;
+    TF_RETURN_IF_ERROR(NotSupportedInImplicitBatch());
+
+    const auto& inputs = params.inputs;
+    const auto& node_def = params.node_def;
+    const TRT_TensorOrWeights& dims_input = inputs.at(0);
+
+    const auto dims_type = dims_input.TrtDType();
+    if (dims_type != nvinfer1::DataType::kINT32) {
+      return errors::InvalidArgument("The dims parameter of ", node_def.op(),
+                                     " operation in ", node_def.name(),
+                                     " is expected to be of type ",
+                                     DebugString(nvinfer1::DataType::kINT32),
+                                     " type, got ", DebugString(dims_type));
+    }
+
+    const auto nbDims = dims_input.GetTrtDims().nbDims;
+    if (nbDims < 0) {
+      return errors::InvalidArgument("The shape of parameter ", node_def.op(),
+                                     " operation in ", node_def.name(),
+                                     " cannot be partial.");
+    }
+    return OkStatus();
+  }
+
+  Status Convert() {
+    const auto& params = *this->params_;
+    auto* network = params.converter->network();
+    const auto& inputs = params.inputs;
+
+    const bool is_dims_static = inputs[0].is_weights();
+    const bool is_value_static = inputs[1].is_weights();
+
+    const TRT_TensorOrWeights& dims_input = inputs.at(0);
+    const TRT_TensorOrWeights& value_input = inputs.at(1);
+
+    int nbDims = dims_input.GetTrtDims().d[0];
+
+    nvinfer1::Dims trt_dims{0};
+    if (is_dims_static) {
+      const auto dims_weights = dims_input.weights();
+      DimsAdapter dims_adapter(dims_weights.GetSpan<int32>());
+      dims_adapter.TrtDims(&trt_dims);
+    }
+
+    auto builder = TRTNetworkBuilder::Create(network, params.weight_store);
+    StatusOr<nvinfer1::ILayer*> layer =
+        builder->AddFill(value_input, dims_input, is_value_static,
+                         is_dims_static, nbDims, trt_dims);
+    ITensorProxyPtr output_tensor = (*layer)->getOutput(0);
+    this->AddOutput(TRT_TensorOrWeights(output_tensor));
+    return OkStatus();
+  }
+};
+
+class ConvertRange : public ConvertFillBase<ConvertRange> {
+ public:
+  explicit ConvertRange(const OpConverterParams* params)
+      : ConvertFillBase<ConvertRange>(params) {}
+
+  static constexpr std::array<InputArgSpec, 3> InputSpec() {
+    return std::array<InputArgSpec, 3>{
+        InputArgSpec::Create("start", TrtInputArg::kBoth),
+        InputArgSpec::Create("limit", TrtInputArg::kBoth),
+        InputArgSpec::Create("delta", TrtInputArg::kBoth)};
+  }
+
+  static constexpr const char* NodeDefDataTypeAttributeName() {
+    /*
+    node {
+      name: "..."
+      op: "Range"
+      ...
+      attr {
+        key: "Tidx"
+        value {
+          type: DT_INT32
+        }
+      }
+    }
+    */
+    return "Tidx";
+  }
+  Status Validate() {
+    TF_RETURN_IF_ERROR(NotSupportedInImplicitBatch());
+    const auto& params = *this->params_;
+    const auto& inputs = params.inputs;
+    const auto& node_def = params.node_def;
+
+    float param[3];
+    all_weights_ = all_integers_ = true;
+    for (int i = 0; i < 3; i++) {
+      const auto& input = inputs.at(i);
+      all_integers_ &= input.TrtDType() == nvinfer1::DataType::kINT32;
+      if (input.is_weights()) {
+        switch (input.TrtDType()) {
+          case nvinfer1::DataType::kFLOAT:
+            param[i] = get_input_param<float>(input);
+            break;
+          case nvinfer1::DataType::kHALF:
+            param[i] = get_input_param<Eigen::half>(input);
+            break;
+          case nvinfer1::DataType::kINT32:
+            param[i] = get_input_param<int>(input);
+            break;
+          default:
+            return errors::InvalidArgument(
+                "Unsupported data type ", DebugString(input.TrtDType()),
+                " used for '", InputSpec()[i].name, "'");
+        }
+      } else {
+        all_weights_ = false;
+      }
+    }
+
+    if (!(all_weights_ || all_integers_)) {
+      // As of 06/03/2022, when at least one of the (start, limit, delta)
+      // is passed as a tensor, they must all be of type kINT32
+      return errors::Unimplemented(convert_range_expected_msg(node_def));
+    }
+
+    if (inputs.at(2).is_weights()) {
+      if ((delta_ = param[2]) == 0) {
+        return errors::InvalidArgument("The delta parameter of ", node_def.op(),
+                                       " operation cannot be equal to 0");
+      }
+
+      if (!all_weights_ && delta_ < 0) {
+        return errors::InvalidArgument(
+            "The delta parameter of Range operation "
+            "cannot be negative, when one of (start, limit) is passed as "
+            "a tensor, but got ",
+            delta_);
+      }
+    }
+
+    for (int i = 0; i < 3; i++) {
+      const auto& input = inputs.at(i);
+      const auto& dims = input.GetTrtDims();
+      if (dims.nbDims != 1 || dims.d[0] != 1) {
+        return errors::InvalidArgument("Dimension for '", InputSpec()[i].name,
+                                       "' of ", node_def.op(), " operator ",
+                                       "should be equal to 1");
+      }
+    }
+
+    if (all_weights_) {
+      const auto num_intervals_float =
+          (param[1] - (start_ = param[0])) / delta_;
+      if (num_intervals_float < 0) {
+        const auto error = convert_range_error_msg(start_, param[1], delta_);
+        return errors::InvalidArgument(error);
+      }
+
+      num_values_ = static_cast<int>(num_intervals_float);
+      if (start_ + delta_ * num_values_ != param[1]) {
+        num_values_++;
+      }
+    }
+
+    return OkStatus();
+  }
+
+  Status Convert() {
+    const auto& params = *this->params_;
+    const auto& inputs = params.inputs;
+    const TRT_TensorOrWeights& input = inputs.at(0);
+    TRT_TensorOrWeights value_input;
+    nvinfer1::Dims trt_dims{1};
+    auto builder = TRTNetworkBuilder::Create(params.converter->network(),
+                                             params.weight_store);
+    TRT_ENSURE_OK(builder);
+    ITensorProxyPtr dims_input_tensor = nullptr;
+    ITensorProxyPtr beta_tensor = nullptr;
+    ITensorProxyPtr scalar_tensor = nullptr;
+    if (!all_weights_) {
+      ITensorProxyPtr tensors[3];
+      for (int i = 0; i < 3; i++) {
+        TF_RETURN_IF_ERROR(
+            builder->get_tensor4TensorOrWeights(inputs.at(i), tensors + i));
+      }
+
+      StatusOr<nvinfer1::IElementWiseLayer*> num =
+          builder->Sub(/*limit*/ tensors[1]->trt_tensor(),
+                       /*start*/ tensors[0]->trt_tensor());
+
+      TRT_ENSURE_PTR_OK(num);
+      StatusOr<nvinfer1::IElementWiseLayer*> ceil_div = builder->FloorDiv(
+          (*num)->getOutput(0), (beta_tensor = tensors[2])->trt_tensor());
+      TRT_ENSURE_PTR_OK(ceil_div);
+      dims_input_tensor = (*ceil_div)->getOutput(0);
+      dims_input_tensor->setType(nvinfer1::DataType::kINT32);
+
+      nvinfer1::Dims scalar_dims{0};
+      TF_RETURN_IF_ERROR(PrepareTensorForShape(
+          params.converter, params.inputs.at(0), scalar_dims, false,
+          &scalar_tensor, params.node_def));
+    } else {
+      DimsAdapter value_input_dims(std::vector<int>{1});
+      StatusOr<TRT_ShapedWeights> value_weights =
+          params.weight_store->GetTempWeights(input.TrtDType(),
+                                              value_input_dims);
+
+      TF_RETURN_IF_ERROR(value_weights.status());
+      TF_RETURN_IF_ERROR(value_weights->SetValues(start_));
+      value_input = TRT_TensorOrWeights(value_weights.value());
+
+      trt_dims.d[0] = num_values_;
+      StatusOr<nvinfer1::IConstantLayer*> const_layer =
+          builder->ConstantShape(value_input_dims);
+      TRT_ENSURE_PTR_OK(const_layer);
+      dims_input_tensor = (*const_layer)->getOutput(0);
+    }
+
+    TRT_TensorOrWeights dims_input(dims_input_tensor);
+
+    StatusOr<nvinfer1::ILayer*> layer =
+        builder->AddFill(value_input, dims_input, all_weights_, all_weights_, 1,
+                         trt_dims, scalar_tensor, beta_tensor, delta_);
+
+    ITensorProxyPtr output_tensor = (*layer)->getOutput(0);
+    if (all_integers_) {
+      output_tensor->setType(nvinfer1::DataType::kINT32);
+    }
+
+    this->AddOutput(TRT_TensorOrWeights(output_tensor));
+    return OkStatus();
+  }
+
+ private:
+  template <typename T>
+  float get_input_param(const TRT_TensorOrWeights& input) {
+    return static_cast<float>(*input.weights().GetPointer<T>());
+  }
+
+  float start_;
+  float delta_;
+  int num_values_;
+  bool all_weights_;
+  bool all_integers_;
+};
+
+std::string convert_range_error_msg(float start, float limit, float delta) {
+  constexpr const char* format_string =
+      "For parameters (start, limit) = (%.2f, %.2f) "
+      "of the Range operation delta cannot be %s, got %.2f";
+  return absl::StrFormat(format_string, start, limit,
+                         start < limit ? "negative" : "positive", delta);
+}
+
+std::string convert_range_expected_msg(const NodeDef& node_def) {
+  return "When at least one of parameters (start, limit, delta) of " +
+         node_def.op() + " operation in " + node_def.name() +
+         " is passed as a tensor, they must all be of type kINT32";
+}
+
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertFill>(), "Fill");
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertRange>(),
+                                  "Range");
+
+#endif  // IS_TRT_VERSION_GE(8, 2, 0, 0)
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h

@@ -0,0 +1,715 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_OPS_LAYER_UTILS_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_OPS_LAYER_UTILS_H_
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include <type_traits>
+
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/statusor.h"
+#include "third_party/tensorrt/NvInfer.h"
+#include "third_party/tensorrt/NvInferRuntimeCommon.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+namespace convert {
+
+// Facilitates the creation of TensorRT layers inside a network. The user
+// provides a INetworkDefinition pointer during construction. They can then add
+// operations to the network through the provided functions. Each function
+// returns a struct which contains the symbolic result of the operation (ITensor
+// pointer) as well as a pointer to the last TensorRT ILayer created. Some
+// operations may create multiple layers in order to accomplish the desired
+// result (e.g. Sign).
+class TRTNetworkBuilder {
+ public:
+  static StatusOr<TRTNetworkBuilder> Create(
+      nvinfer1::INetworkDefinition* network, TrtWeightStore* weight_store) {
+    TRT_ENSURE(network);
+    TRT_ENSURE(weight_store);
+    return TRTNetworkBuilder(network, weight_store);
+  }
+
+ private:
+  TRTNetworkBuilder(nvinfer1::INetworkDefinition* network,
+                    TrtWeightStore* weight_store)
+      : network_(network), weight_store_(weight_store) {}
+
+ public:
+  // Adds an Add operation to the network.
+  StatusOr<nvinfer1::IElementWiseLayer*> Add(nvinfer1::ITensor* lhs,
+                                             nvinfer1::ITensor* rhs) noexcept {
+    TRT_ENSURE(lhs);
+    TRT_ENSURE(rhs);
+    nvinfer1::IElementWiseLayer* layer = network_->addElementWise(
+        *lhs, *rhs, nvinfer1::ElementWiseOperation::kSUM);
+    TRT_ENSURE(layer);
+    return layer;
+  };
+
+  // Adds an elementwise min(lhs, rhs) operation to the network. The output has
+  // the same data type as the input.
+  StatusOr<nvinfer1::IElementWiseLayer*> Min(nvinfer1::ITensor* lhs,
+                                             nvinfer1::ITensor* rhs) noexcept {
+    TRT_ENSURE(lhs);
+    TRT_ENSURE(rhs);
+    nvinfer1::IElementWiseLayer* layer = network_->addElementWise(
+        *lhs, *rhs, nvinfer1::ElementWiseOperation::kMIN);
+    TRT_ENSURE(layer);
+    return layer;
+  };
+
+  // Adds an elementwise max(lhs, rhs) operation to the network. The output has
+  // the same datatype as the input.
+  StatusOr<nvinfer1::IElementWiseLayer*> Max(nvinfer1::ITensor* lhs,
+                                             nvinfer1::ITensor* rhs) noexcept {
+    TRT_ENSURE(lhs);
+    TRT_ENSURE(rhs);
+    nvinfer1::IElementWiseLayer* layer = network_->addElementWise(
+        *lhs, *rhs, nvinfer1::ElementWiseOperation::kMAX);
+    TRT_ENSURE(layer);
+    return layer;
+  };
+
+  // Adds an absolute value operation to the network. Note that this unary
+  // operation will do an implicit float conversion. For int32 tensors, use
+  // "AbsInt".
+  StatusOr<nvinfer1::IUnaryLayer*> AbsFloat(nvinfer1::ITensor* input) noexcept {
+    TRT_ENSURE(input);
+    TRT_ENSURE(input->getType() != nvinfer1::DataType::kFLOAT &&
+               input->getType() != nvinfer1::DataType::kHALF);
+    nvinfer1::IUnaryLayer* layer =
+        network_->addUnary(*input, nvinfer1::UnaryOperation::kABS);
+    TRT_ENSURE(layer);
+    return layer;
+  }
+
+  // Performs Abs without implicit float conversion. The input should be of type
+  // kInt32. For float datatypes, use "Abs".
+  StatusOr<nvinfer1::IElementWiseLayer*> AbsInt(
+      nvinfer1::ITensor* input) noexcept {
+    TRT_ENSURE(input);
+    TRT_ENSURE(input->getType() == nvinfer1::DataType::kINT32);
+    StatusOr<nvinfer1::IElementWiseLayer*> sign = this->SignInt(input);
+    return this->Mul(input, (*sign)->getOutput(0));
+  }
+
+  // Returns elementwise sign(x) for int32 input tensors where sign(x) is
+  // defined as 1 where x > 0, -1 where x < 0 and 0 where x == 0.
+  StatusOr<nvinfer1::IElementWiseLayer*> SignInt(
+      nvinfer1::ITensor* input) noexcept {
+    TRT_ENSURE(input);
+
+    // Create constants +1 and -1.
+    StatusOr<nvinfer1::IConstantLayer*> one =
+        this->Constant<int32>(1, input->getDimensions().nbDims);
+    TRT_ENSURE_PTR_OK(one);
+
+    StatusOr<nvinfer1::IConstantLayer*> neg_one =
+        this->Constant<int32>(-1, input->getDimensions().nbDims);
+    TRT_ENSURE_PTR_OK(neg_one);
+
+    // Turn all negaitve elements into -1, positive and zero elements
+    // unaffected.
+    StatusOr<nvinfer1::IElementWiseLayer*> max =
+        this->Max(input, (*neg_one)->getOutput(0));
+    TRT_ENSURE_PTR_OK(max);
+
+    // Turn all positive elements into +1, negative and zero elements
+    // unaffected.
+    StatusOr<nvinfer1::IElementWiseLayer*> min =
+        this->Min((*max)->getOutput(0), (*one)->getOutput(0));
+    TRT_ENSURE_PTR_OK(min);
+    return min;
+  }
+
+  // Adds a Sub operation to the network.
+  StatusOr<nvinfer1::IElementWiseLayer*> Sub(nvinfer1::ITensor* lhs,
+                                             nvinfer1::ITensor* rhs) noexcept {
+    TRT_ENSURE(lhs);
+    TRT_ENSURE(rhs);
+    nvinfer1::IElementWiseLayer* layer = network_->addElementWise(
+        *lhs, *rhs, nvinfer1::ElementWiseOperation::kSUB);
+    TRT_ENSURE(layer);
+    return layer;
+  }
+
+  // Adds an Greater operation to the network.
+  StatusOr<nvinfer1::IElementWiseLayer*> Greater(
+      nvinfer1::ITensor* lhs, nvinfer1::ITensor* rhs) noexcept {
+    TRT_ENSURE(lhs);
+    TRT_ENSURE(rhs);
+    nvinfer1::IElementWiseLayer* layer = network_->addElementWise(
+        *lhs, *rhs, nvinfer1::ElementWiseOperation::kGREATER);
+    TRT_ENSURE(layer);
+    return layer;
+  }
+
+  // Adds an Equal operation to the network.
+  StatusOr<nvinfer1::IElementWiseLayer*> Equal(
+      nvinfer1::ITensor* lhs, nvinfer1::ITensor* rhs) noexcept {
+    TRT_ENSURE(lhs);
+    TRT_ENSURE(rhs);
+    nvinfer1::IElementWiseLayer* layer = network_->addElementWise(
+        *lhs, *rhs, nvinfer1::ElementWiseOperation::kEQUAL);
+    TRT_ENSURE(layer);
+    return layer;
+  }
+
+  // Adds a FloorDiv operation to the network.
+  StatusOr<nvinfer1::IElementWiseLayer*> FloorDiv(
+      nvinfer1::ITensor* lhs, nvinfer1::ITensor* rhs) noexcept {
+    TRT_ENSURE(lhs);
+    TRT_ENSURE(rhs);
+    nvinfer1::IElementWiseLayer* layer = network_->addElementWise(
+        *lhs, *rhs, nvinfer1::ElementWiseOperation::kFLOOR_DIV);
+    TRT_ENSURE(layer);
+    return layer;
+  }
+
+  // Returns the equivalent of ceil_divide(abs(x)/abs(y))) operation. The inputs
+  // "lhs" and "rhs" should be int32 tensors.
+  StatusOr<nvinfer1::IElementWiseLayer*> AbsCeilDivInt(
+      nvinfer1::ITensor* lhs, nvinfer1::ITensor* rhs) noexcept {
+    TRT_ENSURE(lhs);
+    TRT_ENSURE(rhs);
+    TRT_ENSURE(lhs->getType() == nvinfer1::DataType::kINT32);
+    TRT_ENSURE(rhs->getType() == nvinfer1::DataType::kINT32);
+
+    StatusOr<nvinfer1::IElementWiseLayer*> rhs_abs = this->AbsInt(rhs);
+    TRT_ENSURE_PTR_OK(rhs_abs);
+    StatusOr<nvinfer1::IElementWiseLayer*> lhs_abs = this->AbsInt(lhs);
+    TRT_ENSURE_PTR_OK(lhs_abs);
+    StatusOr<nvinfer1::IElementWiseLayer*> add1 =
+        this->Add((*lhs_abs)->getOutput(0), (*rhs_abs)->getOutput(0));
+    TRT_ENSURE_PTR_OK(add1);
+    StatusOr<nvinfer1::IConstantLayer*> one_const =
+        this->Constant<int32>(1, rhs->getDimensions().nbDims);
+    TRT_ENSURE_PTR_OK(one_const);
+    StatusOr<nvinfer1::IElementWiseLayer*> numerator =
+        this->Sub((*add1)->getOutput(0), (*one_const)->getOutput(0));
+    TRT_ENSURE_PTR_OK(numerator);
+    return FloorDiv((*numerator)->getOutput(0), (*rhs_abs)->getOutput(0));
+  }
+
+  // Adds an elementwise multiplication operation to the network.
+  StatusOr<nvinfer1::IElementWiseLayer*> Mul(nvinfer1::ITensor* lhs,
+                                             nvinfer1::ITensor* rhs) noexcept {
+    TRT_ENSURE(lhs);
+    TRT_ENSURE(rhs);
+    nvinfer1::IElementWiseLayer* layer = network_->addElementWise(
+        *lhs, *rhs, nvinfer1::ElementWiseOperation::kPROD);
+    TRT_ENSURE(layer);
+    return layer;
+  }
+
+  // Adds a sequence of elementwise multiplication operations to the network.
+  // The returned layer's output contains the cumulative elementwise product of
+  // all tensors in the input.
+  StatusOr<nvinfer1::ILayer*> CumulativeProd(
+      absl::Span<nvinfer1::ITensor*> inputs) noexcept {
+    TRT_ENSURE(!absl::c_any_of(
+        inputs, [](nvinfer1::ITensor* x) { return x == nullptr; }));
+    nvinfer1::ILayer* out = nullptr;
+    if (inputs.size() == 1) {
+      out = network_->addIdentity(*inputs[0]);
+      TRT_ENSURE(out != nullptr);
+      return out;
+    }
+    nvinfer1::ITensor* last = inputs[0];
+    for (int i = 1; i < inputs.size(); i++) {
+      StatusOr<nvinfer1::IElementWiseLayer*> mul = this->Mul(last, inputs[i]);
+      TRT_ENSURE_PTR_OK(mul);
+      out = *mul;
+      last = (*mul)->getOutput(0);
+    }
+    return out;
+  }
+
+  // Adds a Constant layer whose output is a TensorRT shape tensor. The shape
+  // tensor's size and values correspond to dim's nbDims and d[], respectively.
+  StatusOr<nvinfer1::IConstantLayer*> ConstantShape(
+      const DimsAdapter& shape_data) noexcept {
+    TRT_ENSURE(shape_data.NumDims() > 0);
+    nvinfer1::Dims shape_dims;
+    shape_dims.nbDims = 1;
+    shape_dims.d[0] = shape_data.NumDims();
+    StatusOr<TRT_ShapedWeights> const_weights =
+        weight_store_->GetTempWeights(nvinfer1::DataType::kINT32, shape_dims);
+    TRT_ENSURE_OK(const_weights);
+    absl::c_copy(shape_data, const_weights->GetPointer<int32>());
+    StatusOr<nvinfer1::Dims> trt_dims = const_weights->Shape().AsTrtDims();
+    TRT_ENSURE_OK(trt_dims);
+    nvinfer1::IConstantLayer* const_layer =
+        network_->addConstant(*trt_dims, const_weights->GetTrtWeights());
+    TRT_ENSURE(const_layer);
+    nvinfer1::ITensor* output = const_layer->getOutput(0);
+    TRT_ENSURE(output);
+    TRT_ENSURE(output->getType() == nvinfer1::DataType::kINT32);
+    return const_layer;
+  }
+
+  // Adds a Constant layer whose output is a TensorRT shape tensor. The shape
+  // tensor's size and values correspond to dim's nbDims and d[], respectively.
+  StatusOr<nvinfer1::IConstantLayer*> Constant(
+      const std::vector<int>& data) noexcept {
+    nvinfer1::Dims shape_dims;
+    shape_dims.nbDims = 1;
+    shape_dims.d[0] = data.size();
+    StatusOr<TRT_ShapedWeights> const_weights =
+        weight_store_->GetTempWeights(nvinfer1::DataType::kINT32, shape_dims);
+    TRT_ENSURE_OK(const_weights);
+    int32* values = const_weights->GetPointer<int32>();
+    for (int i = 0; i < data.size(); i++) {
+      values[i] = static_cast<int32>(data[i]);
+    }
+    StatusOr<nvinfer1::Dims> trt_dims = const_weights->Shape().AsTrtDims();
+    TRT_ENSURE_OK(trt_dims);
+    nvinfer1::IConstantLayer* const_layer =
+        network_->addConstant(*trt_dims, const_weights->GetTrtWeights());
+    TRT_ENSURE(const_layer);
+    nvinfer1::ITensor* output = const_layer->getOutput(0);
+    TRT_ENSURE(output);
+    TRT_ENSURE(output->getType() == nvinfer1::DataType::kINT32);
+    TRT_ENSURE(const_layer);
+    return const_layer;
+  }
+
+  // Adds a Constant layer that produces a tensor of shape "shape",
+  // type "data_type" and filled with value "scalar".
+  template <typename T>
+  StatusOr<nvinfer1::IConstantLayer*> Constant(
+      const T value, nvinfer1::Dims shape,
+      nvinfer1::DataType data_type) noexcept {
+    StatusOr<TRT_ShapedWeights> const_weights =
+        weight_store_->GetTempWeights(data_type, shape);
+    TRT_ENSURE_OK(const_weights);
+    TRT_ENSURE(const_weights->SetValues(value).ok());
+    nvinfer1::IConstantLayer* const_layer =
+        network_->addConstant(shape, const_weights->GetTrtWeights());
+    TRT_ENSURE(const_layer);
+    return const_layer;
+  }
+
+  // Adds a Constant layer that produces a tensor with a single value "scalar".
+  // The tensor has "nb_dims" dimensions and each dimension has only one
+  // element. The data type of the tensor is determined by the data type of
+  // "scalar".
+  template <typename T, typename std::enable_if<std::is_trivially_copyable<
+                            T>::value>::type* = nullptr>
+  StatusOr<nvinfer1::IConstantLayer*> Constant(const T scalar,
+                                               const int nb_dims) noexcept {
+    TRT_ENSURE(nb_dims <= nvinfer1::Dims::MAX_DIMS);
+    auto data_type = nvinfer1::DataType::kINT32;
+    if (std::is_floating_point<T>::value) {
+      data_type = nvinfer1::DataType::kFLOAT;
+    }
+    nvinfer1::Dims zero_shape;
+    zero_shape.nbDims = nb_dims;
+    std::fill_n(zero_shape.d, nb_dims, 1);
+    return Constant<T>(scalar, zero_shape, data_type);
+  }
+
+  // Adds a Constant layer from a TRT_ShapedWeights object.
+  StatusOr<nvinfer1::IConstantLayer*> WeightsToConstant(
+      const nvinfer1::Weights& weights, const DimsAdapter& dims) noexcept {
+    StatusOr<int64_t> vol = dims.Volume();
+    TRT_ENSURE_OK(vol);
+    TRT_ENSURE(*vol == weights.count);
+    StatusOr<nvinfer1::Dims> trt_dims = dims.AsTrtDims();
+    TRT_ENSURE_OK(trt_dims);
+    nvinfer1::IConstantLayer* const_layer =
+        network_->addConstant(*trt_dims, weights);
+    TRT_ENSURE(const_layer);
+    return const_layer;
+  }
+
+  Status get_tensor4TensorOrWeights(const TRT_TensorOrWeights& input,
+                                    ITensorProxyPtr* pTensor) {
+    if (input.is_weights()) {
+      StatusOr<nvinfer1::IConstantLayer*> const_layer = WeightsToConstant(
+          input.weights().GetTrtWeights(), input.GetTrtDims());
+      if (!const_layer.status().ok()) return const_layer.status();
+      *pTensor = (*const_layer)->getOutput(0);
+    } else {
+      *pTensor = input.tensor();
+    }
+    return OkStatus();
+  }
+
+  // Creates a nvinfer1::Weights object containing a single scalar.
+  template <typename T, typename std::enable_if<std::is_trivially_copyable<
+                            T>::value>::type* = nullptr>
+  StatusOr<nvinfer1::Weights> ScalarWeights(const T scalar,
+                                            const int nb_dims) noexcept {
+    TRT_ENSURE(nb_dims <= nvinfer1::Dims::MAX_DIMS);
+    auto data_type = nvinfer1::DataType::kINT32;
+    if (std::is_floating_point<T>::value) {
+      data_type = nvinfer1::DataType::kFLOAT;
+    }
+    nvinfer1::Dims weights_shape;
+    weights_shape.nbDims = nb_dims;
+    std::fill_n(weights_shape.d, nb_dims, 1);
+    StatusOr<TRT_ShapedWeights> const_weights =
+        weight_store_->GetTempWeights(data_type, weights_shape);
+    TRT_ENSURE_OK(const_weights);
+    const_weights->GetPointer<T>()[0] = scalar;
+    return const_weights->GetTrtWeights();
+  }
+
+  // Adds a TensorRT Slice operation to the network.
+  StatusOr<nvinfer1::ISliceLayer*> Slice(
+      nvinfer1::ITensor* input, const nvinfer1::Dims& begin,
+      const nvinfer1::Dims& size, const nvinfer1::Dims& stride) noexcept {
+    nvinfer1::ISliceLayer* layer =
+        network_->addSlice(*input, begin, size, stride);
+    TRT_ENSURE(layer);
+    return layer;
+  }
+
+  // Adds a TensorRT Concatenate operation to the network.
+  StatusOr<nvinfer1::IConcatenationLayer*> Concat(
+      absl::Span<nvinfer1::ITensor* const> inputs, const int axis) {
+    for (nvinfer1::ITensor* input : inputs) {
+      TRT_ENSURE(input);
+    }
+    nvinfer1::IConcatenationLayer* layer = network_->addConcatenation(
+        inputs.data(), static_cast<int32_t>(inputs.size()));
+    TRT_ENSURE(layer);
+    layer->setAxis(axis);
+    return layer;
+  }
+
+  // Adds a TensorRT Concatenate operation to the network.
+  StatusOr<nvinfer1::IConcatenationLayer*> Concat(
+      const std::vector<nvinfer1::ITensor*>& inputs, const int axis) {
+    return this->Concat(absl::MakeSpan(inputs), axis);
+  }
+
+  // Adds a TensorRT Shape operation, which determines the runtime shape of the
+  // input tensor, to the network.
+  StatusOr<nvinfer1::IShapeLayer*> Shape(nvinfer1::ITensor* input) {
+    TRT_ENSURE(input);
+    nvinfer1::IShapeLayer* layer = network_->addShape(*input);
+    TRT_ENSURE(layer);
+    return layer;
+  }
+
+  // Creates a Gather operation on the shape of the input tensor. The output of
+  // the gather operation is a 1D shape tensor where output[i] = (!sub_one ?
+  // input_shape[i] : input_shape[i] -1) if i is in "indices", otherwise zero.
+  StatusOr<nvinfer1::IGatherLayer*> GetPartialShapeOf(
+      nvinfer1::ITensor* input, absl::InlinedVector<int64, 4> indices,
+      bool sub_one = false) {
+    TRT_ENSURE(input);
+    TRT_ENSURE(indices.size() <= nvinfer1::Dims::MAX_DIMS);
+
+    // Get the runtime shape of input;
+    StatusOr<nvinfer1::IShapeLayer*> shape_layer = this->Shape(input);
+    TRT_ENSURE_PTR_OK(shape_layer);
+    nvinfer1::ITensor* runtime_shape = (*shape_layer)->getOutput(0);
+
+    if (sub_one) {
+      StatusOr<nvinfer1::IConstantLayer*> ones = this->Constant<int32>(1, 1);
+      TRT_ENSURE_PTR_OK(ones);
+      StatusOr<nvinfer1::IElementWiseLayer*> sub =
+          this->Sub(runtime_shape, (*ones)->getOutput(0));
+      TRT_ENSURE_PTR_OK(sub);
+      runtime_shape = (*sub)->getOutput(0);
+    }
+
+    // Create a constant tensor containing the gather indices.
+    // For any dim not in "indices", we mark it size to gather a zero.
+    const int input_nb_dims = input->getDimensions().nbDims;
+    std::vector<int> indices_all(input_nb_dims, input_nb_dims);
+    for (auto idx : indices) {
+      TRT_ENSURE(idx < input_nb_dims);
+      indices_all[idx] = idx;
+    }
+
+    StatusOr<nvinfer1::IConstantLayer*> indices_result =
+        this->Constant(indices_all);
+    TRT_ENSURE_PTR_OK(indices_result);
+    nvinfer1::ITensor* gather_indices = (*indices_result)->getOutput(0);
+    TRT_ENSURE(gather_indices->getDimensions().nbDims == 1);
+    TRT_ENSURE(gather_indices->getType() == nvinfer1::DataType::kINT32);
+
+    // Append a zero to the shape tensor.
+    StatusOr<nvinfer1::IConstantLayer*> zero_result =
+        this->Constant(std::vector<int>{0});
+    TRT_ENSURE_PTR_OK(zero_result);
+    std::array<nvinfer1::ITensor*, 2> cat_inputs = {
+        runtime_shape, (*zero_result)->getOutput(0)};
+    nvinfer1::IConcatenationLayer* cat_layer =
+        network_->addConcatenation(cat_inputs.data(), cat_inputs.size());
+    TRT_ENSURE(cat_layer);
+    nvinfer1::ITensor* gather_input = cat_layer->getOutput(0);
+    TRT_ENSURE(gather_input);
+
+    // Finally, gather the indices from the input.
+    nvinfer1::IGatherLayer* gather =
+        network_->addGather(*gather_input, *gather_indices, 0);
+    TRT_ENSURE(gather);
+    return gather;
+  }
+
+  // Adds a scale layer that uniformly scales the input tensor by the specified
+  // amount.
+  StatusOr<nvinfer1::IScaleLayer*> AddUniformScale(nvinfer1::ITensor* input,
+                                                   float scale,
+                                                   const std::string& name) {
+    TRT_ENSURE(input);
+    TRT_ENSURE(!name.empty());
+    StatusOr<nvinfer1::Weights> weight = this->ScalarWeights<float>(scale, 1);
+    TRT_ENSURE_OK(weight);
+    const nvinfer1::Weights empty_weights =
+        nvinfer1::Weights{nvinfer1::DataType::kFLOAT, nullptr, 0};
+    nvinfer1::IScaleLayer* scale_layer =
+        network_->addScale(*input, nvinfer1::ScaleMode::kUNIFORM, empty_weights,
+                           (*weight), empty_weights);
+    TRT_ENSURE(scale_layer != nullptr);
+    scale_layer->setName(name.c_str());
+    TRT_ENSURE((*scale_layer).getPower().count == 0);
+    TRT_ENSURE((*scale_layer).getShift().count == 0);
+    TRT_ENSURE((*scale_layer).getScale().count == 1);
+    return scale_layer;
+  }
+
+  StatusOr<nvinfer1::ILayer*> AddFill(const TRT_TensorOrWeights& value_input,
+                                      const TRT_TensorOrWeights& dims_input,
+                                      bool is_value_static, bool is_dims_static,
+                                      int nbDims,
+                                      const nvinfer1::Dims& trt_dims,
+                                      ITensorProxyPtr scalar_tensor = nullptr,
+                                      ITensorProxyPtr beta_tensor = nullptr,
+                                      const float delta = 0) {
+    // TensorRT IFillLayer requires a rank 0 scalar.
+    nvinfer1::Dims scalar_dims;
+    scalar_dims.nbDims = 0;
+    if (is_value_static) {
+      StatusOr<nvinfer1::IConstantLayer*> const_layer =
+          WeightsToConstant(value_input.weights().GetTrtWeights(), scalar_dims);
+      if (!const_layer.status().ok()) return const_layer.status();
+      scalar_tensor = (*const_layer)->getOutput(0);
+    } else {
+      if (scalar_tensor == nullptr) {
+        StatusOr<nvinfer1::IShuffleLayer*> shuffler_layer =
+            Reshape(value_input.tensor()->trt_tensor(), scalar_dims);
+        if (!shuffler_layer.status().ok()) return shuffler_layer.status();
+        scalar_tensor = (*shuffler_layer)->getOutput(0);
+      }
+    }
+
+    if (beta_tensor == nullptr) {
+      nvinfer1::Dims beta_shape{1, {nbDims}};
+      StatusOr<nvinfer1::IConstantLayer*> const_layer =
+          Constant(delta, beta_shape, value_input.TrtDType());
+      TF_RETURN_IF_ERROR(const_layer.status());
+      beta_tensor = (*const_layer)->getOutput(0);
+    }
+
+    nvinfer1::IFillLayer* layer =
+        network_->addFill(trt_dims, nvinfer1::FillOperation::kLINSPACE);
+    TRT_ENSURE(layer);
+    if (!is_dims_static) {
+      layer->setInput(0, *dims_input.tensor()->trt_tensor());
+    }
+    layer->setInput(1, *scalar_tensor->trt_tensor());
+    layer->setInput(2, *beta_tensor->trt_tensor());
+    return layer;
+  }
+
+  // Adds a quantization layer that uniformly scales the input tensor
+  // by the given multiplicative "scaling_factor", then rounds
+  // (round-to-nearest-ties-to-even) to the nearest integer and clamps in the
+  // range of [-128, 127].
+  StatusOr<nvinfer1::ILayer*> Quantize(nvinfer1::ITensor* input,
+                                       const float scaling_factor,
+                                       const std::string& name) {
+    TRT_ENSURE(input);
+    TRT_ENSURE(!name.empty());
+    // Preprocessor usage here is unavoidable because TRT8 API is new.
+#if IS_TRT_VERSION_GE(8, 0, 0, 0)
+    // The TensorRT IQuantizeLayer divides by the scale factor rather than
+    // multiplies. To be consistent, in this function we expect a multiplicative
+    // scale factor, so we take the reciprical.
+    StatusOr<nvinfer1::IConstantLayer*> scaling_const =
+        this->Constant<float>(1.0f / scaling_factor, 1);
+    TRT_ENSURE_PTR_OK(scaling_const);
+    (*scaling_const)->setDimensions(nvinfer1::Dims{0, {}});
+    nvinfer1::IQuantizeLayer* quant_layer =
+        network_->addQuantize(*input, *(*scaling_const)->getOutput(0));
+    TRT_ENSURE(quant_layer);
+    quant_layer->setAxis(1);
+    return quant_layer;
+#else
+    StatusOr<nvinfer1::IScaleLayer*> result =
+        this->AddUniformScale(input, scaling_factor, name);
+    TRT_ENSURE_PTR_OK(result);
+    (*result)->setOutputType(0, nvinfer1::DataType::kINT8);
+    (*result)->setPrecision(nvinfer1::DataType::kFLOAT);
+    return result;
+#endif
+  }
+
+  // Adds a dequantize layer that casts the input tensor to TensorRT float type
+  // and scales it uniformly by the given multiplicative "scaling_factor".
+  StatusOr<nvinfer1::ILayer*> Dequantize(nvinfer1::ITensor* input,
+                                         const float scaling_factor,
+                                         const std::string& name) {
+    TRT_ENSURE(input);
+    TRT_ENSURE(!name.empty());
+#if IS_TRT_VERSION_GE(8, 0, 0, 0)
+    StatusOr<nvinfer1::IConstantLayer*> scaling_const =
+        this->Constant<float>(scaling_factor, 1);
+    TRT_ENSURE_PTR_OK(scaling_const);
+    (*scaling_const)->setDimensions(nvinfer1::Dims{0, {}});
+    nvinfer1::IDequantizeLayer* dequant_layer =
+        network_->addDequantize(*input, *(*scaling_const)->getOutput(0));
+    dequant_layer->setAxis(1);
+    TRT_ENSURE(dequant_layer);
+    return dequant_layer;
+#else
+    StatusOr<nvinfer1::IScaleLayer*> result =
+        this->AddUniformScale(input, scaling_factor, name);
+    TRT_ENSURE_PTR_OK(result);
+    (*result)->setOutputType(0, nvinfer1::DataType::kFLOAT);
+    (*result)->setPrecision(nvinfer1::DataType::kINT8);
+    return result;
+#endif
+  }
+
+  // Adds TensorRT Q/DQ operations. This is for explicit precision mode.
+  StatusOr<nvinfer1::ILayer*> UniformQuantizeDequantizeExplicit(
+      nvinfer1::ITensor* input, float quantize_scale, float dequantize_scale,
+      const std::string& name) {
+    TRT_ENSURE(input);
+    if (!IS_TRT_VERSION_GE(8, 0, 0, 0)) {
+      TRT_ENSURE(network_->hasExplicitPrecision());
+    }
+    TRT_ENSURE(IS_TRT_VERSION_GE(7, 1, 0, 0));
+
+    static int count = 0;
+    TRT_ENSURE(input->getType() == nvinfer1::DataType::kFLOAT);
+    std::string quant_name = absl::StrCat(input->getName(), "_quant_", count);
+
+    StatusOr<nvinfer1::ILayer*> quant =
+        this->Quantize(input, quantize_scale, quant_name);
+    TRT_ENSURE_PTR_OK(quant);
+
+    std::string dequant_name =
+        absl::StrCat(input->getName(), "_dequant_", count);
+    StatusOr<nvinfer1::ILayer*> dequant = this->Dequantize(
+        (*quant)->getOutput(0), dequantize_scale, dequant_name);
+    TRT_ENSURE_PTR_OK(dequant);
+
+    count++;
+    return dequant;
+  }
+
+  StatusOr<nvinfer1::IShuffleLayer*> Reshape(nvinfer1::ITensor* input,
+                                             const nvinfer1::Dims& new_shape) {
+    TRT_ENSURE(input);
+    nvinfer1::IShuffleLayer* layer = network_->addShuffle(*input);
+    TRT_ENSURE(layer);
+    layer->setReshapeDimensions(new_shape);
+    return layer;
+  }
+
+  StatusOr<nvinfer1::ILayer*> FindProducerOf(const nvinfer1::ITensor* tensor) {
+    const char* name = tensor->getName();
+    const int num_layers = network_->getNbLayers();
+    for (int i = 0; i < num_layers; i++) {
+      nvinfer1::ILayer* layer = network_->getLayer(i);
+      const int num_outputs = layer->getNbOutputs();
+      for (int j = 0; j < num_outputs; j++) {
+        nvinfer1::ITensor* t = layer->getOutput(j);
+        if (std::string(t->getName()) == name) {
+          return layer;
+        }
+      }
+    }
+    return errors::NotFound("could not find producing layer of ", name);
+  }
+
+  StatusOr<nvinfer1::ILayer*> UniqueParentOf(const nvinfer1::ILayer* layer,
+                                             int input_idx = 0) {
+    return FindProducerOf(layer->getInput(input_idx));
+  }
+
+  nvinfer1::INetworkDefinition* Network() { return network_; }
+
+ private:
+  nvinfer1::INetworkDefinition* network_;
+  TrtWeightStore* weight_store_;
+};
+
+class ShuffleBuilder {
+ private:
+  explicit ShuffleBuilder(TRTNetworkBuilder* builder, nvinfer1::ITensor* input)
+      : builder_(builder) {
+    layer_ = builder->Network()->addShuffle(*input);
+  }
+
+ public:
+  static StatusOr<ShuffleBuilder> Create(TRTNetworkBuilder* builder,
+                                         nvinfer1::ITensor* input) {
+    TRT_ENSURE(builder != nullptr);
+    TRT_ENSURE(input != nullptr);
+    return ShuffleBuilder(builder, input);
+  }
+
+  ShuffleBuilder& SetReshape(const nvinfer1::Dims& dims) {
+    layer_->setReshapeDimensions(dims);
+    return *this;
+  }
+
+  ShuffleBuilder& SetReshape(nvinfer1::ITensor* shape) {
+    layer_->setInput(1, *shape);
+    return *this;
+  }
+
+  ShuffleBuilder& SetFirstTranspose(const nvinfer1::Permutation& perm) {
+    layer_->setFirstTranspose(perm);
+    return *this;
+  }
+
+  ShuffleBuilder& SetSecondTranspose(const nvinfer1::Permutation& perm) {
+    layer_->setSecondTranspose(perm);
+    return *this;
+  }
+
+  StatusOr<nvinfer1::ITensor*> Output() {
+    TRT_ENSURE(layer_ != nullptr);
+    TRT_ENSURE(layer_->getOutput(0) != nullptr);
+    return layer_->getOutput(0);
+  }
+
+ private:
+  TRTNetworkBuilder* builder_;
+  nvinfer1::IShuffleLayer* layer_;
+};
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_OPS_LAYER_UTILS_H_

tensorflow/compiler/tf2tensorrt/convert/ops/like_ops.cc

@@ -0,0 +1,94 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+#if IS_TRT_VERSION_GE(8, 2, 0, 0)
+
+template <int V>
+class ConvertLikeOps : public OpConverterBase<ConvertLikeOps<V>> {
+ public:
+  explicit ConvertLikeOps(const OpConverterParams *params)
+      : OpConverterBase<ConvertLikeOps<V>>(
+            params,
+            {DataType::DT_FLOAT, DataType::DT_HALF, DataType::DT_INT32}) {}
+
+  static constexpr std::array<InputArgSpec, 1> InputSpec() {
+    return std::array<InputArgSpec, 1>{
+        InputArgSpec::Create("input", TrtInputArg::kBoth),
+    };
+  }
+  Status Validate() { return ConvertLikeOps<V>::NotSupportedInImplicitBatch(); }
+
+  Status Convert() {
+    const auto &params = *this->params_;
+    const auto &inputs = params.inputs;
+    auto *network = params.converter->network();
+
+    const TRT_TensorOrWeights &input = inputs.at(0);
+    nvinfer1::Dims dims(input.GetTrtDims());
+
+    const std::vector<int> value_input_dims_data = {1};
+    const DimsAdapter value_input_dims(value_input_dims_data);
+    StatusOr<TRT_ShapedWeights> value_weights =
+        params.weight_store->GetTempWeights(input.TrtDType(), value_input_dims);
+    TF_RETURN_IF_ERROR(value_weights.status());
+    TF_RETURN_IF_ERROR(value_weights->SetValues(V));
+    TRT_TensorOrWeights value_input(value_weights.value());
+
+    const auto is_dims_static = HasStaticShape(dims);
+    auto builder = TRTNetworkBuilder::Create(network, params.weight_store);
+    ITensorProxyPtr dims_input_tensor;
+    if (!is_dims_static) {
+      StatusOr<nvinfer1::IShapeLayer *> shape_layer =
+          builder->Shape(input.tensor()->trt_tensor());
+      TF_RETURN_IF_ERROR(shape_layer.status());
+      dims_input_tensor = (*shape_layer)->getOutput(0);
+      dims.nbDims = 0;
+    }
+
+    TRT_TensorOrWeights dims_input(dims_input_tensor);
+    StatusOr<nvinfer1::ILayer *> layer =
+        builder->AddFill(value_input, dims_input, true, is_dims_static,
+                         input.GetTrtDims().nbDims, dims);
+    ITensorProxyPtr output_tensor = (*layer)->getOutput(0);
+    this->AddOutput(TRT_TensorOrWeights(output_tensor));
+    return OkStatus();
+  }
+};
+
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertLikeOps<0>>(),
+                                  "zeros_like");
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertLikeOps<1>>(),
+                                  "ones_like");
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertLikeOps<0>>(),
+                                  "ZerosLike");
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertLikeOps<1>>(),
+                                  "OnesLike");
+
+#endif  // IS_TRT_VERSION_GE(8, 2, 0, 0)
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/log_softmax.cc

@@ -0,0 +1,104 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+class ConvertLogSoftmax : public OpConverterBase<ConvertLogSoftmax> {
+ public:
+  explicit ConvertLogSoftmax(const OpConverterParams *params)
+      : OpConverterBase<ConvertLogSoftmax>(params) {}
+
+  static constexpr std::array<InputArgSpec, 1> InputSpec() {
+    return std::array<InputArgSpec, 1>{
+        InputArgSpec::Create("logits", TrtInputArg::kTensor)};
+  }
+
+  Status Validate() {
+    const auto &params = *this->params_;
+    const auto &inputs = params.inputs;
+
+    ITensorProxyPtr logits_tensor = inputs.at(0).tensor();
+
+    const int num_trt_dims = logits_tensor->getDimensions().nbDims;
+    if (!num_trt_dims && params.use_implicit_batch) {
+      return errors::InvalidArgument(
+          "TensorRT LogSoftmax cannot apply on the batch dimension");
+    }
+
+    return OkStatus();
+  }
+
+  Status Convert() {
+    const auto &params = *this->params_;
+    const auto &inputs = params.inputs;
+    const auto &node_def = params.node_def;
+
+    // Perform LogSoftmax operation:
+    // `logsoftmax = logits - log(reduce_sum(exp(logits), axis))`
+
+    // Get the logits tensor.
+    ITensorProxyPtr logits_tensor = inputs.at(0).tensor();
+    const int num_trt_dims = logits_tensor->getDimensions().nbDims;
+
+    // Exponent of logits.
+    nvinfer1::IUnaryLayer *exp = params.converter->network()->addUnary(
+        *logits_tensor->trt_tensor(), nvinfer1::UnaryOperation::kEXP);
+    TFTRT_RETURN_ERROR_IF_NULLPTR(exp, node_def.name());
+    params.converter->SetLayerName(exp, node_def, "exp");
+
+    // Reduce-sum operation across the final dimension.
+    nvinfer1::IReduceLayer *reduced_sum =
+        params.converter->network()->addReduce(
+            *exp->getOutput(0), nvinfer1::ReduceOperation::kSUM,
+            (1 << (num_trt_dims - 1)), /*Reduce across final dimension*/
+            true /*Keep reduced dims*/);
+    params.converter->SetLayerName(reduced_sum, node_def, "reduced_sum");
+
+    // Logarithm of reduced_sum.
+    nvinfer1::IUnaryLayer *log_reduced_sum =
+        params.converter->network()->addUnary(*reduced_sum->getOutput(0),
+                                              nvinfer1::UnaryOperation::kLOG);
+    TFTRT_RETURN_ERROR_IF_NULLPTR(log_reduced_sum, node_def.name());
+    params.converter->SetLayerName(log_reduced_sum, node_def,
+                                   "log_reduced_sum");
+
+    // Finally, get the output by subtracting log_reduced_sum from logits.
+    nvinfer1::IElementWiseLayer *sub =
+        params.converter->network()->addElementWise(
+            *logits_tensor->trt_tensor(), *log_reduced_sum->getOutput(0),
+            nvinfer1::ElementWiseOperation::kSUB);
+    TFTRT_RETURN_ERROR_IF_NULLPTR(sub, node_def.name());
+    params.converter->SetLayerName(sub, node_def, "sub");
+
+    params.outputs->push_back(TRT_TensorOrWeights(sub->getOutput(0)));
+    return OkStatus();
+  }
+};
+
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertLogSoftmax>(),
+                                  "LogSoftmax");
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/quantization_ops.cc

@@ -0,0 +1,421 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/quantization_ops.h"
+
+#include "absl/strings/str_format.h"
+#include "tensorflow/cc/ops//array_ops.h"
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/weights.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "third_party/tensorrt/NvInfer.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+bool IsQuantizeAndDequantizeOp(const Node* node) {
+  return absl::c_find(kQuantizationOpNames, node->def().op()) !=
+         kQuantizationOpNames.end();
+}
+
+namespace {
+
+// Provides quantizing and dequantizing tensor scales for a given dynamic range.
+// Borrowed from TF quantization kernel logic.
+template <typename T>
+QuantizationScales<T, 1> ComputeQuantizationRange(bool signed_input,
+                                                  int num_bits,
+                                                  bool narrow_range,
+                                                  T* min_range, T* max_range) {
+  // Calculate the range for the simulated integer quantization:
+  // e.g. [-127,127] for signed = true, narrow_range = true, num_bits = 8,
+  // or [-128,127] for signed = true, narrow_range = false, num_bits = 8,
+  // or [0, 255] for signed = false, num_bits = 8.
+  const int64_t min_quantized =
+      signed_input ? narrow_range ? -(1ULL << (num_bits - 1)) + 1
+                                  : -(1ULL << (num_bits - 1))
+                   : 0;
+  const int64_t max_quantized =
+      signed_input ? (1ULL << (num_bits - 1)) - 1 : (1ULL << num_bits) - 1;
+  // Determine the maximum scaling factor that would scale
+  // [min_range, max_range] to not exceed [min_quantized, max_quantized],
+  // while keeping 0 unchanged.
+  const T scale_from_min_side = (min_quantized * *min_range > 0)
+                                    ? min_quantized / *min_range
+                                    : std::numeric_limits<T>::max();
+  const T scale_from_max_side = (max_quantized * *max_range > 0)
+                                    ? max_quantized / *max_range
+                                    : std::numeric_limits<T>::max();
+
+  QuantizationScales<T, 1> scales;
+  // Note: Avoids changing the side of the range that determines scale.
+  if (scale_from_min_side < scale_from_max_side) {
+    scales.quantize_scale[0] = scale_from_min_side;
+    scales.dequantize_scale[0] = *min_range / min_quantized;
+    *max_range = max_quantized * scales.dequantize_scale[0];
+  } else {
+    scales.quantize_scale[0] = scale_from_max_side;
+    scales.dequantize_scale[0] = *max_range / max_quantized;
+    *min_range = min_quantized * scales.dequantize_scale[0];
+  }
+  return scales;
+}
+
+// Prepares the input for a QDQ node in explicit precision mode, returning a
+// ITensor pointer. If the input is weights, we convert it to a ITensor by
+// adding a constant layer.
+StatusOr<nvinfer1::ITensor*> ExlicitQDQInputToTensor(
+    TRTNetworkBuilder* builder, const OpConverterParams* params,
+    const TRT_TensorOrWeights& input) {
+  if (input.is_tensor()) {
+    return input.tensor()->trt_tensor();
+  }
+  if (!IS_TRT_VERSION_GE(8, 0, 0, 0) && input.weights().count() > 1) {
+    LOG(WARNING) << absl::StrCat(
+        "QDQ per-channel for weights not "
+        "implemented, assuming uniform scaling");
+  }
+  TRT_ShapedWeights trt_weights = input.weights();
+  StatusOr<nvinfer1::IConstantLayer*> weights_const =
+      builder->WeightsToConstant(trt_weights.GetTrtWeights(),
+                                 trt_weights.Shape());
+  TRT_ENSURE_PTR_OK(weights_const);
+  params->converter->SetLayerName(*weights_const, params->node_def, "const");
+  nvinfer1::ITensor* qdq_input = (*weights_const)->getOutput(0);
+  std::string name = absl::StrCat((*weights_const)->getName(), "_output");
+  qdq_input->setName(name.c_str());
+  return qdq_input;
+}
+
+}  // namespace
+
+// Carries traits for each specific quantization op type for conversion.
+// Specialization for template parameter T should be given for each TF C++
+// quantization op.
+template <typename T>
+struct QDQOpSpec {};
+
+template <>
+struct QDQOpSpec<ops::QuantizeAndDequantizeV2> {
+  static constexpr std::array<InputArgSpec, 3> InputSpec() {
+    return {
+        InputArgSpec::Create("input", TrtInputArg::kBoth),
+        InputArgSpec::Create("input_min", TrtInputArg::kWeight),
+        InputArgSpec::Create("input_max", TrtInputArg::kWeight),
+    };
+  }
+
+  struct Attrs {
+    float min_range;
+    float max_range;
+    bool narrow_range;
+    std::string round_mode;
+    UniformQuantizationScales scales;
+  };
+
+  static Status ValidateQDQForExplicitPrecision(
+      const std::vector<TRT_TensorOrWeights>& inputs, const NodeDef& node_def,
+      Attrs* args) {
+    AttrSlice attrs(node_def);
+    TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "round_mode", &args->round_mode));
+    if (args->round_mode != "HALF_TO_EVEN") {
+      LOG(WARNING) << node_def.op() << ": " << node_def.name()
+                   << " has round_mode=" << args->round_mode
+                   << ", but for TensorRT conversion, "
+                      "round_mode=HALF_TO_EVEN is recommended.";
+    }
+    TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "narrow_range", &args->narrow_range));
+    if (args->narrow_range) {
+      LOG(WARNING) << node_def.op() << ": " << node_def.name()
+                   << " has narrow_range=true, but for TensorRT conversion, "
+                      "narrow_range=false is recommended.";
+    }
+    args->min_range = inputs.at(1).weights().template GetPointer<float>()[0];
+    args->max_range = inputs.at(2).weights().template GetPointer<float>()[0];
+    const int num_bits = 8;
+    args->scales = ComputeQuantizationRange<float>(
+        /*signed_input=*/true, num_bits, args->narrow_range, &args->min_range,
+        &args->max_range);
+    TRT_ENSURE(args->scales.dequantize_scale[0] != 0);
+    TRT_ENSURE(args->scales.quantize_scale[0] != 0);
+    return OkStatus();
+  }
+
+  // Converts in explicit precision mode. In this mode, QDQ operations are
+  // directly converted into TensorRT quantizing and dequantizing scale
+  // operations.
+  static Status ConvertExplicit(const OpConverterParams* params,
+                                const Attrs& args) {
+    const auto& node_def = params->node_def;
+
+    StatusOr<TRTNetworkBuilder> builder = TRTNetworkBuilder::Create(
+        params->converter->network(), params->weight_store);
+
+    StatusOr<nvinfer1::ITensor*> qdq_input =
+        ExlicitQDQInputToTensor(&*builder, params, params->inputs.at(0));
+    TRT_ENSURE_PTR_OK(qdq_input);
+
+    // TODO(cbate): check this condition exists for TRT8? Outline this block to
+    // a "reshape policy".
+    const int required_dims = params->use_implicit_batch ? 3 : 4;
+    const nvinfer1::Dims idims = (*qdq_input)->getDimensions();
+    nvinfer1::Dims intermediate_dims = idims;
+    TRT_ENSURE(idims.nbDims > 0);
+    if (idims.nbDims < required_dims) {
+      const int nb_extra_dims = required_dims - idims.nbDims;
+      intermediate_dims.nbDims = required_dims;
+      std::vector<int> ones(nb_extra_dims, 1);
+      TRT_ENSURE(ones.size() == nb_extra_dims && nb_extra_dims > 0);
+
+      if (!params->use_implicit_batch) {
+        intermediate_dims.d[0] = idims.d[0];
+        std::copy(ones.begin(), ones.end(), intermediate_dims.d + 1);
+        std::copy_n(idims.d + 1, idims.nbDims - 1,
+                    intermediate_dims.d + ones.size() + 1);
+      } else {
+        std::copy(ones.begin(), ones.end(), intermediate_dims.d);
+        std::copy_n(idims.d, idims.nbDims, intermediate_dims.d + ones.size());
+      }
+
+      LOG(WARNING) << absl::StrCat(
+          node_def.name(), ":", node_def.op(), ": tensor ",
+          (*qdq_input)->getName(), " has shape ", DebugString(idims),
+          " but TRT scale layer requires at least 3 dims excluding batch dim, "
+          "trying to recover by inserting 1's to create shape ",
+          DebugString(intermediate_dims));
+      StatusOr<nvinfer1::IShuffleLayer*> reshape =
+          builder->Reshape(*qdq_input, intermediate_dims);
+      TRT_ENSURE_PTR_OK(reshape);
+      *qdq_input = (*reshape)->getOutput(0);
+    }
+
+    VLOG(1) << "[ExplicitPrecision]" << node_def.op() << ": " << node_def.name()
+            << " computed scales: " << args.scales << " from min/max ranges "
+            << args.min_range << "/" << args.max_range;
+
+    StatusOr<nvinfer1::ILayer*> qdq =
+        builder->UniformQuantizeDequantizeExplicit(
+            *qdq_input, args.scales.quantize_scale[0],
+            args.scales.dequantize_scale[0], node_def.name());
+    TRT_ENSURE_PTR_OK(qdq);
+    ITensorProxyPtr final_output = (*qdq)->getOutput(0);
+    if (idims.nbDims != intermediate_dims.nbDims) {
+      StatusOr<nvinfer1::IShuffleLayer*> undo_reshape =
+          builder->Reshape(*qdq_input, idims);
+      TRT_ENSURE_PTR_OK(undo_reshape);
+      final_output = (*undo_reshape)->getOutput(0);
+    }
+    params->outputs->push_back(final_output);
+    return OkStatus();
+  }
+};
+
+template <>
+
+struct QDQOpSpec<ops::QuantizeAndDequantizeV3> {
+  static constexpr std::array<InputArgSpec, 4> InputSpec() {
+    return {
+        InputArgSpec::Create("input", TrtInputArg::kBoth),
+        InputArgSpec::Create("min", TrtInputArg::kWeight),
+        InputArgSpec::Create("max", TrtInputArg::kWeight),
+        InputArgSpec::Create("num_bits", TrtInputArg::kWeight),
+    };
+  }
+  // Use same attributes and conversion functions as QDQV2.
+  using Attrs = QDQOpSpec<ops::QuantizeAndDequantizeV2>::Attrs;
+
+  static Status ValidateQDQForExplicitPrecision(
+      const std::vector<TRT_TensorOrWeights>& inputs, const NodeDef& node_def,
+      Attrs* args) {
+    return QDQOpSpec<
+        ops::QuantizeAndDequantizeV2>::ValidateQDQForExplicitPrecision(inputs,
+                                                                       node_def,
+                                                                       args);
+  }
+
+  static Status ConvertExplicit(const OpConverterParams* params,
+                                const Attrs& args) {
+    return QDQOpSpec<ops::QuantizeAndDequantizeV2>::ConvertExplicit(params,
+                                                                    args);
+  }
+};
+
+template <>
+
+struct QDQOpSpec<ops::FakeQuantWithMinMaxVars> {
+  static constexpr std::array<InputArgSpec, 3> InputSpec() {
+    return {
+        InputArgSpec::Create("input", TrtInputArg::kBoth),
+        InputArgSpec::Create("min", TrtInputArg::kWeight),
+        InputArgSpec::Create("max", TrtInputArg::kWeight),
+    };
+  }
+  struct Attrs {
+    int num_bits;
+    bool narrow_range;
+  };
+
+  static Status ValidateQDQForExplicitPrecision(
+      const std::vector<TRT_TensorOrWeights>& inputs, const NodeDef& node_def,
+      Attrs* args) {
+    return errors::Unimplemented("");
+  }
+
+  static Status ConvertExplicit(const OpConverterParams* params,
+                                const Attrs& args) {
+    return errors::Unimplemented("");
+  }
+};
+
+template <>
+
+struct QDQOpSpec<ops::FakeQuantWithMinMaxArgs> {
+  static constexpr std::array<InputArgSpec, 1> InputSpec() {
+    return {
+        InputArgSpec::Create("input", TrtInputArg::kBoth),
+    };
+  }
+
+  struct Attrs {
+    float min;
+    float max;
+    int num_bits;
+    bool narrow_range;
+  };
+
+  static Status ValidateQDQForExplicitPrecision(
+      const std::vector<TRT_TensorOrWeights>& inputs, const NodeDef& node_def,
+      Attrs* args) {
+    return errors::Unimplemented("");
+  }
+
+  static Status ConvertExplicit(const OpConverterParams* params,
+                                const Attrs& args) {
+    return errors::Unimplemented("");
+  }
+};
+
+// Converts QDQ operations in non-explicit precision mode. This is the original
+// "ConvertQuantize" function. In this mode, Q/DQ operations are no-ops and are
+// instead used to set the dynamic range of the input tensor.
+Status ConvertDynamicRangeMode(const OpConverterParams* params) {
+  const auto& inputs = params->inputs;
+  const auto& node_def = params->node_def;
+  float min_range = 0.0f;
+  float max_range = 0.0f;
+  const auto& op_name = node_def.op();
+  if (op_name == "FakeQuantWithMinMaxArgs") {
+    AttrSlice attrs(node_def);
+    // Get ranges via node attributes.
+    TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "min", &min_range));
+    TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "max", &max_range));
+  } else if (op_name == "FakeQuantWithMinMaxVars" ||
+             op_name == "QuantizeAndDequantizeV2" ||
+             op_name == "QuantizeAndDequantizeV3") {
+    // Get ranges via inputs.
+    auto get_weights_value = [&inputs](int index) {
+      const auto* raw_weights = inputs.at(index).weights().GetPointer<float>();
+      return raw_weights[0];
+    };
+    min_range = get_weights_value(1);
+    max_range = get_weights_value(2);
+  } else {
+    return errors::InvalidArgument("Unknown quantization op ", op_name, ", at ",
+                                   node_def.name());
+  }
+  if (params->validation_only) {
+    return OkStatus();
+  }
+
+  // Store ranges for tensor
+  ITensorProxyPtr input0 = inputs.at(0).tensor();
+  params->converter->ProvideQuantizationRange(&input0, min_range, max_range);
+  // Sometimes, TRT may not quantize a tensor, either because it chooses to
+  // execute a higher precision kernel or because of op fusion. In these
+  // cases, accuracy will suffer if the model was trained to expect
+  // quantization at that tensor. We should consider adding a clip(tensor,
+  // min_range, max_range) operation here to ensure that any arbitrarily
+  // placed quantize node will execute as expected. However, this will
+  // negatively affect performance. If users train their models in a way which
+  // models inference as close as possible (i.e. not quantizing in place where
+  // fusion will occur), then there is no problem with the current
+  // implementation.
+  params->outputs->push_back(inputs.at(0));
+  return OkStatus();
+}
+
+template <typename TFOpType>
+class ConvertQDQ : public OpConverterBase<ConvertQDQ<TFOpType>> {
+ public:
+  explicit ConvertQDQ(const OpConverterParams* params)
+      : OpConverterBase<ConvertQDQ<TFOpType>>(params) {}
+
+  static constexpr auto InputSpec() { return QDQOpSpec<TFOpType>::InputSpec(); }
+
+  // Disable the non-applicable data type check by providing empty string.
+  static constexpr const char* NodeDefDataTypeAttributeName() { return ""; }
+
+  Status ValidateDynamicRangeINT8Mode() {
+    // The condition ensures we only call the conversion once. We should break
+    // this function up into validation and conversion.
+    if (this->params_->validation_only) {
+      return ConvertDynamicRangeMode(this->params_);
+    }
+    return OkStatus();
+  }
+
+  Status Validate() {
+    if (!this->params_->use_explicit_precision) {
+      return ValidateDynamicRangeINT8Mode();
+    }
+    return OpSpec::ValidateQDQForExplicitPrecision(
+        this->params_->inputs, this->params_->node_def, &attrs_);
+  }
+
+  Status Convert() {
+    if (!this->params_->use_explicit_precision) {
+      return ConvertDynamicRangeMode(this->params_);
+    }
+    return OpSpec::ConvertExplicit(this->params_, attrs_);
+  }
+
+  using OpSpec = QDQOpSpec<TFOpType>;
+  using OpSpecAttrs = typename QDQOpSpec<TFOpType>::Attrs;
+  OpSpecAttrs attrs_;
+};
+
+REGISTER_DEFAULT_TRT_OP_CONVERTER(
+    MakeConverterFunction<ConvertQDQ<ops::QuantizeAndDequantizeV2>>(),
+    "QuantizeAndDequantizeV2");
+REGISTER_DEFAULT_TRT_OP_CONVERTER(
+    MakeConverterFunction<ConvertQDQ<ops::QuantizeAndDequantizeV3>>(),
+    "QuantizeAndDequantizeV3");
+REGISTER_DEFAULT_TRT_OP_CONVERTER(
+    MakeConverterFunction<ConvertQDQ<ops::FakeQuantWithMinMaxVars>>(),
+    "FakeQuantWithMinMaxVars");
+REGISTER_DEFAULT_TRT_OP_CONVERTER(
+    MakeConverterFunction<ConvertQDQ<ops::FakeQuantWithMinMaxArgs>>(),
+    "FakeQuantWithMinMaxArgs");
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/quantization_ops.h

@@ -0,0 +1,76 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_OPS_QUANTIZATION_OPS_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_OPS_QUANTIZATION_OPS_H_
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "absl/strings/str_format.h"
+#include "absl/strings/str_join.h"
+#include "tensorflow/core/graph/graph.h"
+#include "tensorflow/core/lib/core/status.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+constexpr std::array<const char*, 4> kQuantizationOpNames = {
+    "QuantizeAndDequantizeV2",
+    "QuantizeAndDequantizeV3",
+    "FakeQuantWithMinMaxVars",
+    "FakeQuantWithMinMaxArgs",
+};
+
+// Operations with supported conversion to Q/DQ ops in TensorRT explicit
+// precision mode.
+constexpr std::array<const char*, 1> kExplicitQuantizationOpNames = {
+    "QuantizeAndDequantizeV2",
+};
+
+// Contains two scaling factors for quantization and dequantization
+// respectively. A shift factor is omitted as TensorRT only supports symmetric
+// quantization.
+template <typename T, size_t N>
+struct QuantizationScales {
+  std::array<T, N> quantize_scale;
+  std::array<T, N> dequantize_scale;
+};
+
+// In TensorRT 7 and 8, only uniform tensor scaling is supported for
+// activations.
+using UniformQuantizationScales = QuantizationScales<float, 1>;
+
+// Per-channel scaling is supported for weights in TensorRT version >= 8.0.
+template <size_t ChannelDimSize>
+using PerChannelQuantizationScales = QuantizationScales<float, ChannelDimSize>;
+
+template <typename T, size_t N>
+std::ostream& operator<<(std::ostream& os,
+                         const QuantizationScales<T, N>& scales) {
+  os << absl::StrFormat("QuantizationScales[quantize={%s},dequantize={%s}]",
+                        absl::StrJoin(scales.quantize_scale, ","),
+                        absl::StrJoin(scales.dequantize_scale, ","));
+  return os;
+}
+
+// Returns true if the Tensorflow node is a quantize and dequantize operation.
+bool IsQuantizeAndDequantizeOp(const Node*);
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_OPS_QUANTIZATION_OPS_H_

tensorflow/compiler/tf2tensorrt/convert/ops/quantization_ops_test.cc

@@ -0,0 +1,618 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/quantization_ops.h"
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/strings/str_format.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/const_op.h"
+#include "tensorflow/cc/ops/linalg_ops.h"
+#include "tensorflow/cc/ops/math_ops.h"
+#include "tensorflow/cc/ops/nn_ops.h"
+#include "tensorflow/compiler/jit/shape_inference.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/trt_convert_api.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/platform/status_matchers.h"
+#include "tensorflow/core/protobuf/error_codes.pb.h"
+
+#if IS_TRT_VERSION_GE(8, 0, 0, 0)
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+namespace ops = ::tensorflow::ops;
+using ::tensorflow::testing::StatusIs;
+
+// This anonymous namespace contains helper functions for instantiating small TF
+// building blocks. These are used below to construct specific graph patterns
+// which test end-to-end conversion of the TF graph to an explciit-precision
+// enabled TensorRT network.
+namespace {
+
+enum class ConvEpilogueType {
+  kNone,
+  kReLU,
+  kBatchNorm,
+  kReLUBatchnorm,
+  kBatchnormReLU
+};
+
+std::ostream& operator<<(std::ostream& os, ConvEpilogueType epilogue) {
+  switch (epilogue) {
+    case ConvEpilogueType::kNone:
+      return os << "None";
+    case ConvEpilogueType::kReLU:
+      return os << "ReLU only";
+    case ConvEpilogueType::kBatchNorm:
+      return os << "BatchNorm Only";
+    case ConvEpilogueType::kReLUBatchnorm:
+      return os << "ReLU+Batchnorm";
+    case ConvEpilogueType::kBatchnormReLU:
+      return os << "BatchNorm+ReLU";
+  }
+}
+
+std::string DebugString(ConvEpilogueType epilogue) {
+  std::stringstream ss;
+  ss << epilogue;
+  return ss.str();
+}
+
+// Adds a 2D 3x3, single channel input with specified data_format. data_format
+// must be NHWC,NCHW or NHW.
+ops::Placeholder AddInput(Scope scope, int input_idx,
+                          const std::string data_format,
+                          std::array<int, 3> size_chw = {1, 3, 3}) {
+  PartialTensorShape input_shape;
+  if (data_format == "NCHW") {
+    input_shape =
+        PartialTensorShape({1, size_chw[0], size_chw[1], size_chw[2]});
+  } else if (data_format == "NHWC") {
+    input_shape =
+        PartialTensorShape({1, size_chw[1], size_chw[2], size_chw[0]});
+  } else if (data_format == "NHW") {
+    input_shape = PartialTensorShape({1, size_chw[1], size_chw[2]});
+  } else {
+    LOG(FATAL) << "Unknown input shape type " << data_format;
+  }
+  auto input_attrs = ops::Placeholder::Attrs().Shape(input_shape);
+  return ops::Placeholder(scope.WithOpName(absl::StrCat("input_", input_idx)),
+                          DT_FLOAT, input_attrs);
+}
+
+// Adds QDQ op with min = -1.0f, max = 1.0f.
+Output AddQDQV2(Scope scope, Input input) {
+  // Create scaling factors.
+  auto input_min =
+      ops::Const<float>(scope.WithOpName("in_min"), -1.0f, TensorShape{});
+  auto input_max =
+      ops::Const<float>(scope.WithOpName("in_max"), 1.0f, TensorShape{});
+  return ops::QuantizeAndDequantizeV2(scope.WithOpName("qdq"), input, input_min,
+                                      input_max);
+}
+
+Output AddOutput(Scope scope, Output input, int idx, bool add_qdq) {
+  Output out = input;
+  if (add_qdq) {
+    out = AddQDQV2(scope, input);
+  }
+  return ops::Identity(scope.WithOpName(StrCat("output_", idx)), out);
+}
+
+// Adds a 3x3x1x1 Conv2D op and optional bias weights, followed by ReLU
+// activation. Puts QDQ between (weights, op). Puts QDQ between (input, op)
+// when qdq_on_output=false. Otherwise, puts QDQ between (op, output).
+Output AddConv2D(Scope scope, Input input, int in_channels, int out_channels,
+                 std::array<int, 2> filter_size = {1, 1},
+                 std::array<int, 2> stride = {1, 1},
+                 const std::string& data_format = "NCHW", bool with_bias = true,
+                 ConvEpilogueType epilogue = ConvEpilogueType::kBatchnormReLU,
+                 bool qdq_on_output = false) {
+  // Create 3x3 non-quantized weights weights.
+  auto weights_const = ops::Const(
+      scope.WithOpName("weights"), 1.0f,
+      TensorShape({filter_size[0], filter_size[1], in_channels, out_channels}));
+
+  // Add QDQ to input if we don't add QDQ to output.
+  auto conv_input =
+      !qdq_on_output ? AddQDQV2(scope.WithOpName("qdq_input"), input) : input;
+
+  Output result = ops::Conv2D(
+      scope.WithOpName("conv2d"), conv_input, AddQDQV2(scope, weights_const),
+      /*strides=*/{1, 1, 1, 1},
+      /*padding=*/"SAME", ops::Conv2D::Attrs().DataFormat(data_format));
+
+  if (with_bias) {
+    auto bias_const = ops::Const(scope.WithOpName("bias_weights"), 1.0f,
+                                 TensorShape({
+                                     out_channels,
+                                 }));
+    result = ops::BiasAdd(scope.WithOpName("bias"), result, bias_const,
+                          ops::BiasAdd::Attrs().DataFormat(data_format));
+  }
+
+  auto add_bn = [scope, data_format](Input input,
+                                     const int channels) -> Output {
+    TensorShape constant_shape = TensorShape({channels});
+    auto bn_scale =
+        ops::Const(scope.WithOpName("bn_scale"), 1.0f, constant_shape);
+    auto bn_offset =
+        ops::Const(scope.WithOpName("bn_offset"), 1.0f, constant_shape);
+    auto bn_mean =
+        ops::Const(scope.WithOpName("bn_mean"), 0.1f, TensorShape({channels}));
+    auto bn_var =
+        ops::Const(scope.WithOpName("bn_var"), 1.0f, TensorShape({channels}));
+    Input conv_bn_input = IS_TRT_VERSION_GE(8, 0, 1, 0)
+                              ? input
+                              : AddQDQV2(scope.WithOpName("qdq_input"), input);
+    return ops::FusedBatchNormV3(
+               scope.WithOpName("bn"), conv_bn_input, bn_scale, bn_offset,
+               bn_mean, bn_var,
+               ops::FusedBatchNormV3::Attrs().IsTraining(false).DataFormat(
+                   data_format))
+        .y;
+  };
+
+  switch (epilogue) {
+    case ConvEpilogueType::kBatchNorm: {
+      result = add_bn(result, out_channels);
+      break;
+    }
+    case ConvEpilogueType::kReLU: {
+      result = ops::Relu(scope.WithOpName("relu"), result);
+      break;
+    }
+    case ConvEpilogueType::kReLUBatchnorm: {
+      result = ops::Relu(scope.WithOpName("relu"), result);
+      result = add_bn(result, out_channels);
+      break;
+    }
+    case ConvEpilogueType::kBatchnormReLU: {
+      result = add_bn(result, out_channels);
+      result = ops::Relu(scope.WithOpName("relu"), result);
+      break;
+    }
+    case ConvEpilogueType::kNone:
+      break;
+  }
+
+  if (qdq_on_output) {
+    result = AddQDQV2(scope.WithOpName("qdq_out"), result);
+  }
+  return result;
+}
+
+// Adds a batch matrix multiplication V2 operation, which commonly appears in
+// fully connected layers. Puts QDQ between (input, op) as well as between
+// (weights, op).
+ops::BatchMatMulV2 AddMatMul(Scope scope, const std::string& name,
+                             Input input) {
+  // Add QDQ to input.
+  auto input_qdq = AddQDQV2(scope, input);
+
+  // Add 3x3 weights with QDQ.
+  auto weights_const =
+      ops::Const(scope.WithOpName(name + "_weights"),
+                 {1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f},
+                 TensorShape({3, 3}));
+  auto weights_qdq = AddQDQV2(scope.WithOpName("weights_qdq"), weights_const);
+  return ops::BatchMatMulV2(scope.WithOpName(name), input_qdq, weights_qdq);
+}
+}  // namespace
+
+struct QDQTestOptions {
+  bool conv_has_bias{true};
+
+  // TRT7 may have issues with optimizing redundant transpose operations between
+  // QDQ and Op introduced by TF-TRT when format is not "NCHW". This allows to
+  // test both cases as well as WAR feasibility.
+  std::string data_format{"NCHW"};
+
+  // Tests whether placing QDQ on outputs rather than inputs is handled
+  // correctly.
+  bool qdq_on_output{false};
+
+  // Option for testing whether TRT build succeeds without a final QDQ before
+  // the output.
+  bool final_qdq{true};
+
+  // Whether to add activations (relu) to conv operations
+  ConvEpilogueType conv_epilogue;
+
+  // TF-TRT API Options
+  TfTrtConversionParams conversion_params{};
+};
+
+std::ostream& operator<<(std::ostream& os, const QDQTestOptions opts) {
+  return os << absl::StrCat(
+             "QDQTestOptions(conv_has_bias=",
+             static_cast<int>(opts.conv_has_bias),
+             ", qdq_on_output=", static_cast<int>(opts.qdq_on_output),
+             ", data_format=", opts.data_format,
+             ", conv_epilogue=", DebugString(opts.conv_epilogue),
+             ", final_qdq=", opts.final_qdq, ")");
+}
+
+std::vector<QDQTestOptions> EnumerateQDQTestOptions() {
+  std::vector<QDQTestOptions> result;
+  for (const absl::string_view data_format : {"NCHW", "NHWC"}) {
+    for (auto use_bias : {true, false}) {
+      for (auto qdq_on_output : {false, true}) {
+        // For now, always append a QDQ before output. For small single-op tests
+        // (besides QDQ), TensorRT7 sometimes has trouble.
+        for (auto final_qdq : {true, false}) {
+          for (auto conv_epilogue :
+               {ConvEpilogueType::kReLU, ConvEpilogueType::kNone,
+                ConvEpilogueType::kBatchnormReLU}) {
+            // Currently batch norm converter only supports NHWC.
+            if (data_format == "NHWC" &&
+                (conv_epilogue == ConvEpilogueType::kBatchnormReLU ||
+                 conv_epilogue == ConvEpilogueType::kBatchNorm ||
+                 conv_epilogue == ConvEpilogueType::kBatchnormReLU)) {
+              continue;
+            }
+            QDQTestOptions opts{};
+            opts.conv_has_bias = use_bias;
+            opts.data_format = data_format;
+            opts.qdq_on_output = qdq_on_output;
+            opts.final_qdq = final_qdq;
+            opts.conv_epilogue = conv_epilogue;
+            result.push_back(opts);
+          }
+        }
+      }
+    }
+  }
+  return result;
+}
+
+// This class is a test fixture for running graph conversion and evaluating
+// numerical results.
+class QDQExplicitTest : public ::testing::Test,
+                        public ::testing::WithParamInterface<QDQTestOptions> {
+ public:
+  static StatusOr<PartialTensorShape> GetShape(const std::string& name,
+                                               const GraphShapeInfo& shapes) {
+    TRT_ENSURE(shapes.find(name) != shapes.end());
+    TRT_ENSURE(shapes.at(name).size() == 1);
+    return shapes.at(name)[0].shape;
+  }
+
+  StatusOr<MetaGraphDef> GetModel(const GraphDef& graph_def,
+                                  const std::vector<const NodeDef*>& inputs,
+                                  const std::vector<const NodeDef*>& outputs,
+                                  const GraphShapeInfo& shapes) {
+    TRT_ENSURE(!inputs.empty());
+    TRT_ENSURE(!outputs.empty());
+
+    MetaGraphDef out;
+    out.mutable_graph_def()->CopyFrom(graph_def);
+
+    SignatureDef signature_def;
+    auto& mutable_inputs = *signature_def.mutable_inputs();
+    for (int i = 0; i < inputs.size(); i++) {
+      std::string input_name = inputs[i]->name();
+      auto& input = mutable_inputs[input_name];
+      input.set_name(input_name);
+      input.set_dtype(DT_FLOAT);
+      TRT_ENSURE(shapes.find(input_name) != shapes.end());
+      TRT_ENSURE(shapes.at(input_name).size() == 1);
+      PartialTensorShape input_shape = shapes.at(input_name)[0].shape;
+      input_shape.AsProto(input.mutable_tensor_shape());
+    }
+
+    auto& mutable_outputs = *signature_def.mutable_outputs();
+    for (int i = 0; i < outputs.size(); i++) {
+      std::string output_name = outputs[i]->name();
+      auto& output = mutable_outputs[output_name];
+      output.set_name(output_name);
+      output.set_dtype(DT_FLOAT);
+      TRT_ENSURE(shapes.find(output_name) != shapes.end());
+      TRT_ENSURE(shapes.at(output_name).size() == 1);
+      PartialTensorShape output_shape = shapes.at(output_name)[0].shape;
+      output_shape.AsProto(output.mutable_tensor_shape());
+    }
+
+    (*out.mutable_signature_def())["serving_default"] = signature_def;
+    return out;
+  }
+
+  // Confirms that we have a TRT node with the correct attributes.
+  static Status CheckTrtNode(const GraphDef& converted_graph_def) {
+    int n_trt_ops = 0;
+    string op_name{"TRTEngineOp"};
+    for (const auto& node : converted_graph_def.node()) {
+      if (op_name == node.op()) {
+        n_trt_ops++;
+        const auto& attr = node.attr();
+        TRT_ENSURE(attr.at("static_engine").b());
+        VLOG(2) << "Found serialized segment with size "
+                << attr.at("serialized_segment").s().size();
+        TRT_ENSURE(!attr.at("serialized_segment").s().empty());
+      }
+    }
+    TRT_ENSURE(n_trt_ops == 1);
+    return OkStatus();
+  }
+
+  Status ConvertAndRun(Scope* scope) {
+    std::vector<const NodeDef*> inputs;
+    std::vector<const NodeDef*> outputs;
+
+    GraphDef gdef;
+    TF_RETURN_IF_ERROR(scope->ToGraphDef(&gdef));
+
+    std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+    TF_RETURN_IF_ERROR(scope->ToGraph(graph.get()));
+
+    GraphShapeInfo shape_info;
+    TF_RETURN_IF_ERROR(InferShapes(graph.get(), /*arg_shapes=*/{},
+                                   /*fnlib_def=*/nullptr, &shape_info));
+
+    for (const NodeDef& node : gdef.node()) {
+      if (absl::StartsWith(node.name(), "input_")) {
+        inputs.push_back(&node);
+      } else if (absl::StartsWith(node.name(), "output_")) {
+        outputs.push_back(&node);
+      }
+    }
+
+    StatusOr<MetaGraphDef> meta_graph_def =
+        GetModel(gdef, inputs, outputs, shape_info);
+    TRT_ENSURE_OK(meta_graph_def);
+
+    // Create a list of input tensors, they will be used to build the engines.
+    std::vector<Tensor> input_tensors;
+    std::vector<std::string> input_names;
+    for (const auto& input : inputs) {
+      input_names.push_back(input->name());
+
+      StatusOr<PartialTensorShape> input_shape =
+          GetShape(input->name(), shape_info);
+      TRT_ENSURE_OK(input_shape);
+
+      TensorShape shape;
+      input_shape->AsTensorShape(&shape);
+      Tensor tensor(DT_FLOAT, shape);
+      test::FillIota(&tensor, 1.0f);
+      input_tensors.push_back(tensor);
+    }
+
+    std::vector<std::string> output_names;
+    for (const auto& output : outputs) {
+      output_names.push_back(output->name());
+    }
+
+    TfTrtConversionParams conversion_params;
+    conversion_params.allow_build_at_runtime = true;
+    conversion_params.precision_mode = TrtPrecisionMode::INT8;
+    conversion_params.use_calibration = false;
+    conversion_params.convert_to_static_engine = true;
+    TRT_ENSURE(input_names.size() == input_tensors.size());
+    StatusOr<GraphDef> converted_gdef = tensorrt::ConvertAndBuild(
+        meta_graph_def->graph_def(), input_names, output_names, {input_tensors},
+        conversion_params);
+    TRT_ENSURE_OK(converted_gdef);
+    return CheckTrtNode(*converted_gdef);
+  }
+
+ protected:
+  TfTrtConversionParams params_;
+  TrtUniquePtrType<nvinfer1::ICudaEngine> engine_;
+};
+
+class TestQDQSuite : public QDQExplicitTest {};
+
+#define EXPECT_QDQ_ON_OUTPUT_FAILURE(params, scope)                  \
+  if ((params).qdq_on_output) {                                      \
+    EXPECT_THAT(ConvertAndRun(&(scope)), StatusIs(error::INTERNAL)); \
+    return;                                                          \
+  }
+#define EXPECT_NO_FINAL_QDQ_FAILURE(params, scope)                   \
+  if (!(params).final_qdq) {                                         \
+    EXPECT_THAT(ConvertAndRun(&(scope)), StatusIs(error::INTERNAL)); \
+    return;                                                          \
+  }
+
+#define EXPECT_BUILD_OK(scope) TF_EXPECT_OK(ConvertAndRun(&(scope)))
+
+#define POLICY_TRT7(params, scope)               \
+  if (!IS_TRT_VERSION_GE(8, 0, 0, 0)) {          \
+    EXPECT_QDQ_ON_OUTPUT_FAILURE(params, scope); \
+    EXPECT_NO_FINAL_QDQ_FAILURE(params, scope);  \
+    EXPECT_BUILD_OK(scope);                      \
+  }
+
+#define POLICY_TRT8(params, scope)                                          \
+  if (IS_TRT_VERSION_GE(8, 0, 0, 0)) {                                      \
+    if (((params).conv_epilogue == ConvEpilogueType::kBatchNorm ||          \
+         (params).conv_epilogue == ConvEpilogueType::kBatchnormReLU ||      \
+         (params).conv_epilogue == ConvEpilogueType::kReLUBatchnorm) &&     \
+        (params).data_format == "NHWC") {                                   \
+      EXPECT_THAT(ConvertAndRun(&(scope)), StatusIs(error::UNIMPLEMENTED)); \
+      return;                                                               \
+    }                                                                       \
+    EXPECT_BUILD_OK(scope);                                                 \
+  }
+
+#define SKIP_TRT7(x)                           \
+  if (!IS_TRT_VERSION_GE(8, 0, 0, 0) && (x)) { \
+    GTEST_SKIP();                              \
+  }
+
+// Tests single convolution operation conversion.
+TEST_P(TestQDQSuite, TestConv2DBasic) {
+  SKIP_TRT7(GetParam().qdq_on_output);
+  SKIP_TRT7(GetParam().data_format != "NCHW");
+  SKIP_TRT7(!GetParam().final_qdq);
+
+  Scope scope = Scope::NewRootScope();
+  auto input = AddInput(scope, 0, GetParam().data_format, {3, 28, 28});
+
+  Output out = input;
+  const int num_conv = 1;
+  std::array<int, 2> in_channels = {3, 16};
+  std::array<int, 2> out_channels = {16, 32};
+  for (int i = 0; i < num_conv; i++) {
+    out = AddConv2D(scope.WithOpName(absl::StrCat("conv_", i)), out,
+                    in_channels[i], out_channels[i], /*filter_size=*/{3, 3},
+                    /*stride=*/{1, 1}, GetParam().data_format,
+                    GetParam().conv_has_bias, GetParam().conv_epilogue,
+                    GetParam().qdq_on_output);
+  }
+  out = AddOutput(scope, out, 0, GetParam().final_qdq);
+  POLICY_TRT7(GetParam(), scope);
+  POLICY_TRT8(GetParam(), scope);
+}
+
+// Tests single convolution operation conversion.
+TEST_P(TestQDQSuite, TestMatMulBasic) {
+  // Some param's don't apply, so pick one combination and skip otherwise.
+  if (GetParam().data_format != "NCHW" || !GetParam().conv_has_bias ||
+      GetParam().qdq_on_output ||
+      GetParam().conv_epilogue != ConvEpilogueType::kReLU) {
+    GTEST_SKIP();
+  }
+  Scope scope = Scope::NewRootScope();
+  auto input = AddInput(scope, 0, "NHW");
+  auto matmul_op = AddMatMul(scope, "matmul", input);
+  auto out = AddOutput(scope, matmul_op, 0, GetParam().final_qdq);
+
+  TF_EXPECT_OK(ConvertAndRun(&scope));
+}
+
+// A single input goes through two different Conv2D. Outputs of Conv2D are
+// added together, with QQQ on both branches of ADD.
+TEST_P(TestQDQSuite, AddBothBranchesQDQConvSingleInput) {
+  SKIP_TRT7(!GetParam().final_qdq);
+  SKIP_TRT7(GetParam().data_format != "NCHW");
+
+  Scope scope = Scope::NewRootScope();
+  auto input1 = AddInput(scope, 0, GetParam().data_format,
+                         /*size_chw=*/{3, 28, 28});
+
+  auto conv1 =
+      AddConv2D(scope, input1, 3, 16, /*filter_size=*/{3, 3}, /*stride=*/{1, 1},
+                GetParam().data_format, GetParam().conv_has_bias,
+                GetParam().conv_epilogue, GetParam().qdq_on_output);
+
+  auto conv2 =
+      AddConv2D(scope, input1, 3, 16, /*filter_size=*/{3, 3}, /*stride=*/
+                {1, 1}, GetParam().data_format, GetParam().conv_has_bias,
+                GetParam().conv_epilogue, GetParam().qdq_on_output);
+
+  // In the case of "qdq on output", we don't need to add QDQ.
+  auto add =
+      ops::Add(scope.WithOpName("add"),
+               !GetParam().qdq_on_output ? AddQDQV2(scope, conv1) : conv1,
+               !GetParam().qdq_on_output ? AddQDQV2(scope, conv2) : conv2);
+
+  auto conv3 =
+      AddConv2D(scope.WithOpName("conv3"), conv2, 16, 16, {1, 1}, {1, 1},
+                GetParam().data_format, GetParam().conv_has_bias,
+                GetParam().conv_epilogue, GetParam().qdq_on_output);
+
+  auto out =
+      AddOutput(scope.WithOpName("output"), conv3, 0, GetParam().final_qdq);
+
+  POLICY_TRT7(GetParam(), scope);
+  POLICY_TRT8(GetParam(), scope);
+}
+
+// Tests adding a single tensor to itself, with QQQ on both branches of ADD.
+TEST_P(TestQDQSuite, AddBothBranchesQDQMultipleInput) {
+  // TRT7 QDQ optimizer makes single-input restriction.
+  SKIP_TRT7(true);
+
+  Scope scope = Scope::NewRootScope();
+  auto input1 = AddInput(scope, 0, GetParam().data_format);
+  auto input2 = AddInput(scope, 1, GetParam().data_format);
+  auto add =
+      ops::Add(scope.WithOpName("add"),
+               !GetParam().qdq_on_output ? AddQDQV2(scope, input1) : input1,
+               !GetParam().qdq_on_output ? AddQDQV2(scope, input2) : input2);
+  auto output = AddOutput(scope, add, 0, true);
+  TF_EXPECT_OK(ConvertAndRun(&scope));
+}
+
+// Tests Conv-MaxPool combination
+TEST_P(TestQDQSuite, TestConvMaxpool) {
+  SKIP_TRT7(!GetParam().final_qdq);
+  SKIP_TRT7(GetParam().data_format != "NCHW");
+
+  Scope scope = Scope::NewRootScope();
+  auto input = AddInput(scope, 0, GetParam().data_format,
+                        /*size_chw=*/{3, 28, 28});
+  auto conv1 =
+      AddConv2D(scope, input, 3, 16, /*filter_size=*/{3, 3}, /*stride=*/{1, 1},
+                GetParam().data_format, GetParam().conv_has_bias,
+                GetParam().conv_epilogue, GetParam().qdq_on_output);
+  ops::MaxPool maxpool =
+      ops::MaxPool(scope.WithOpName("maxpool"),
+                   AddQDQV2(scope.WithOpName("mp_qdq_in"), conv1), {1, 1, 1, 1},
+                   {1, 1, 1, 1}, "SAME",
+                   ops::MaxPool::Attrs().DataFormat(GetParam().data_format));
+  auto output =
+      AddOutput(scope.WithOpName("output"), maxpool, 0, GetParam().final_qdq);
+  POLICY_TRT7(GetParam(), scope);
+  POLICY_TRT8(GetParam(), scope);
+}
+
+// Tests QDQ(Conv(QDQ(MaxPool(Conv(QDQ(x))))))
+TEST_P(TestQDQSuite, TestConvMaxpoolConv) {
+  SKIP_TRT7(!GetParam().final_qdq);
+  SKIP_TRT7(GetParam().data_format != "NCHW");
+
+  Scope scope = Scope::NewRootScope();
+  auto input = AddInput(scope, 0, GetParam().data_format,
+                        /*size_chw=*/{3, 28, 28});
+  auto conv1 =
+      AddConv2D(scope, input, 3, 16, /*filter_size=*/{3, 3}, /*stride=*/{1, 1},
+                GetParam().data_format, GetParam().conv_has_bias,
+                GetParam().conv_epilogue, GetParam().qdq_on_output);
+  ops::MaxPool maxpool =
+      ops::MaxPool(scope.WithOpName("maxpool"),
+                   AddQDQV2(scope.WithOpName("mp_qdq_in"), conv1), {1, 1, 1, 1},
+                   {1, 1, 1, 1}, "SAME",
+                   ops::MaxPool::Attrs().DataFormat(GetParam().data_format));
+  auto conv2 = AddConv2D(scope, maxpool, 16, 16, {3, 3}, {1, 1},
+                         GetParam().data_format, GetParam().conv_has_bias,
+                         GetParam().conv_epilogue, GetParam().qdq_on_output);
+  auto output =
+      AddOutput(scope.WithOpName("out"), conv2, 0, GetParam().final_qdq);
+  POLICY_TRT7(GetParam(), scope);
+  POLICY_TRT8(GetParam(), scope);
+}
+
+INSTANTIATE_TEST_SUITE_P(TestQDQSuiteInst, TestQDQSuite,
+                         ::testing::ValuesIn(EnumerateQDQTestOptions()));
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // IS_TRT_VERSION_GE(8, 0, 0, 0)
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/selectv2.cc

@@ -0,0 +1,220 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+#if IS_TRT_VERSION_GE(8, 2, 0, 0)
+/*  The ConvertSelectV2 is working only for cond_input passed as a boolean
+ *  tensor, which could be created only for TRT >= 8.2.
+ */
+class ConvertSelectBase : public OpConverterBase<ConvertSelectBase> {
+ public:
+  explicit ConvertSelectBase(const OpConverterParams* params,
+                             const std::string& layer_name)
+      : OpConverterBase<ConvertSelectBase>(
+            params,
+            {DataType::DT_FLOAT, DataType::DT_HALF, DataType::DT_INT32}),
+        layer_name_(layer_name) {}
+
+  static constexpr std::array<InputArgSpec, 3> InputSpec() {
+    return std::array<InputArgSpec, 3>{
+        InputArgSpec::Create("cond", TrtInputArg::kBoth),
+        InputArgSpec::Create("then", TrtInputArg::kBoth),
+        InputArgSpec::Create("else", TrtInputArg::kBoth)};
+  }
+
+  Status Validate() {
+    TF_RETURN_IF_ERROR(NotSupportedInImplicitBatch());
+
+    const auto& params = *this->params_;
+    const auto& inputs = params.inputs;
+    const auto& i_cond = inputs.at(0);
+    const auto& node = params.node_def;
+    TF_RETURN_IF_ERROR(
+        check_type(i_cond.TrtDType(), nvinfer1::DataType::kBOOL, node));
+
+    if (i_cond.is_weights()) {
+      // According to
+      // https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#constant-layer
+      // Boolean weights are not supported in TRT version 8.4.
+      return errors::InvalidArgument(bool_weight_error_msg(node));
+    }
+
+    const auto& i_then = inputs.at(1);
+    const auto& i_else = inputs.at(2);
+    const auto type_then = i_then.TrtDType();
+    const auto type_else = i_else.TrtDType();
+    if (type_then != type_else && (type_then == nvinfer1::DataType::kINT32 ||
+                                   type_else == nvinfer1::DataType::kINT32)) {
+      // Both or none of (type_then, type_else) should be equal to kINT32.
+      return errors::InvalidArgument(
+          then_else_dtypes_error_msg(type_then, type_else, node));
+    }
+
+    bool cond_is_vector = false;
+    const auto& shape_cond = i_cond.GetTrtDims();
+    if (layer_name_ == "select") {
+      const auto& shape_then = i_then.GetTrtDims();
+      const auto& shape_else = i_else.GetTrtDims();
+      TF_RETURN_IF_ERROR(compare_shapes(shape_then, shape_else));
+      TF_RETURN_IF_ERROR(
+          compare_shapes(shape_cond, shape_then, &cond_is_vector));
+    }
+
+    nvinfer1::Dims cond_dims(shape_cond);
+    if (cond_is_vector) {
+      cond_dims.nbDims = i_then.GetTrtDims().nbDims;
+      const std::vector<int> ones(cond_dims.d[0], 1);
+      std::copy(ones.begin(), ones.end(), cond_dims.d + 1);
+    }
+
+    const TRT_TensorOrWeights new_cond(nvinfer1::DataType::kBOOL, cond_dims,
+                                       i_cond.batch_size());
+    nvinfer1::Dims broadcasted_dims[3];
+    for (int i = 1; i < 3; i++) {
+      TF_RETURN_IF_ERROR(GetTrtBroadcastShape(new_cond, inputs.at(i), true,
+                                              false, broadcasted_dims,
+                                              broadcasted_dims + i));
+    }
+
+    for (int i = 0; i < tensor_.size(); i++) {
+      // This will also convert constants to tensors.
+      tensor_[i] = std::make_unique<TRT_TensorOrWeights>(inputs.at(i));
+      TF_RETURN_IF_ERROR(
+          ApplyBroadcast(tensor_[i], broadcasted_dims[i], this->params_, 0));
+    }
+
+    return OkStatus();
+  }
+
+  Status Convert() {
+    const auto& params = *this->params_;
+    auto* converter = params.converter;
+
+    nvinfer1::ISelectLayer* select_layer = converter->network()->addSelect(
+        *tensor_[0].get()->as_tensor(params_)->trt_tensor(),  // cond_tensor
+        *tensor_[1].get()->as_tensor(params_)->trt_tensor(),  // then_tensor
+        *tensor_[2].get()->as_tensor(params_)->trt_tensor()   // else_tensor
+    );
+
+    converter->SetLayerName(select_layer, params.node_def.name(), layer_name_);
+    AddOutput(TRT_TensorOrWeights(select_layer->getOutput(0)));
+    return OkStatus();
+  }
+
+ private:
+  Status compare_shapes(const nvinfer1::Dims& shape1,
+                        const nvinfer1::Dims& shape2,
+                        bool* cond_is_vector = nullptr) const {
+    const bool then_vs_else = cond_is_vector == nullptr;
+    bool same_shapes = shape1 == shape2;
+    if (!same_shapes && shape1.nbDims == shape2.nbDims) {
+      // We can't check size equivalent when dynamic shapes are involved.
+      // In this case, the two shapes should be equal at runtime. Therefore,
+      // the shapes still should be considered as equal if at least one of
+      // them is a tensor with dynamic shape,
+      same_shapes = DynamicShapeInput(this->params_->inputs, then_vs_else);
+    }
+    if (!same_shapes) {
+      if (then_vs_else || !(*cond_is_vector = (shape1.nbDims == 1 &&
+                                               shape1.d[0] == shape2.d[0]))) {
+        const auto err = input_shapes_error_msg(
+            shape1, shape2, this->params_->node_def, then_vs_else);
+        return errors::InvalidArgument(err);
+      }
+    }
+    return OkStatus();
+  }
+
+  bool DynamicShapeInput(const std::vector<TRT_TensorOrWeights>& inputs,
+                         bool then_vs_else) const {
+    const int idx = then_vs_else ? 1 : 0;
+    for (int i = 0; i < 2; ++i) {
+      const auto& input = inputs.at(i + idx);
+      if (input.is_tensor() && !HasStaticShape(input.GetTrtDims())) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  std::array<std::unique_ptr<TRT_TensorOrWeights>, 3> tensor_;
+  const std::string layer_name_;
+};
+
+class ConvertSelect : public ConvertSelectBase {
+ public:
+  explicit ConvertSelect(const OpConverterParams* params)
+      : ConvertSelectBase(params, "select") {}
+};
+
+class ConvertSelectV2 : public ConvertSelectBase {
+ public:
+  explicit ConvertSelectV2(const OpConverterParams* params)
+      : ConvertSelectBase(params, "selectv2") {}
+};
+
+std::string op_node_info(const NodeDef& node) {
+  return " of the '" + node.op() + "' operation at the node '" + node.name() +
+         "' ";
+}
+
+std::string bool_weight_error_msg(const NodeDef& node) {
+  return "The boolean parameter '" + node.input(0) + "'" + op_node_info(node) +
+         "cannot be passed as a weight in TRT version 8.4.";
+}
+
+std::string then_else_dtypes_error_msg(nvinfer1::DataType type_then,
+                                       nvinfer1::DataType type_else,
+                                       const NodeDef& node) {
+  return "DataTypes (" + DebugString(type_then) + ", " +
+         DebugString(type_else) + ") of parameters (" + node.input(1) + ", " +
+         node.input(2) + ")" + op_node_info(node) + "are incompatible.";
+}
+
+std::string input_shapes_error_msg(const nvinfer1::Dims& shape1,
+                                   const nvinfer1::Dims& shape2,
+                                   const NodeDef& node, bool then_vs_else) {
+  const std::string& param_names =
+      then_vs_else ? "'then' and 'else'" : "'cond' and 'then'";
+  std::string error_msg = "The shapes of the " + param_names + " parameters" +
+                          op_node_info(node) + "must be the same";
+  if (!then_vs_else) {
+    error_msg +=
+        " OR 'cond' must be a vector with N elements, "
+        "where N is a batch size (the first shape dimension for 'then')";
+  }
+  return error_msg + ", got " + DebugString(shape1) + " vs. " +
+         DebugString(shape2) + ".";
+}
+
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertSelect>(),
+                                  "Select");
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertSelectV2>(),
+                                  "SelectV2");
+#endif
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/slice_ops.cc

@@ -0,0 +1,300 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/slice_ops.h"
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include <bitset>
+#include <vector>
+
+#include "absl/container/inlined_vector.h"
+#include "absl/strings/str_format.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/optional.h"
+#include "absl/types/span.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h"
+#include "tensorflow/compiler/tf2tensorrt/utils/trt_tensor_proxy.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/util/strided_slice_op.h"
+#include "third_party/tensorrt/NvInfer.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+// Adds a set of operations to the network which set the parameters for the
+// given "slice_layer" in order to handle dynamic input shape.
+Status HandleDynamicStridedSliceInput(
+    TRTNetworkBuilder* builder, nvinfer1::ISliceLayer* slice_layer,
+    const StridedSliceShapeSpec& strided_slice_spec,
+    const absl::InlinedVector<int64, 4>& dynamic_input_size_indices,
+    nvinfer1::Dims begin_dims, nvinfer1::Dims stride_dims,
+    nvinfer1::Dims end_dims);
+
+Status ConvertStridedSliceHelper(
+    const OpConverterParams* params, const TRT_TensorOrWeights& input,
+    const PartialTensorShape& input_dims, const SliceDims& begin,
+    const SliceDims& stride, const SliceDims& end,
+    std::optional<nvinfer1::Dims> final_shape, std::optional<int> op_instance,
+    std::optional<StridedSliceShapeSpec> strided_slice_spec) {
+  const auto& node_def = params->node_def;
+
+  auto begin_dims = DimsAdapter::Create(begin, params->use_implicit_batch);
+  auto stride_dims = DimsAdapter::Create(stride, params->use_implicit_batch);
+  auto end_dims = DimsAdapter::Create(end, params->use_implicit_batch);
+  TRT_ENSURE_OK(begin_dims);
+  TRT_ENSURE_OK(stride_dims);
+  TRT_ENSURE_OK(end_dims);
+
+  // For each dimension, gather information about static vs dynamic dimension
+  // and slice size.
+  nvinfer1::Dims size_dims = begin_dims->AsTrtDims();
+  absl::InlinedVector<int64, 4> static_input_size_indices;
+  absl::InlinedVector<int64, 4> dynamic_input_size_indices;
+  for (int i = 0; i < begin_dims->NumDims(); i++) {
+    size_dims.d[i] = (std::abs(end_dims->dim(i) - begin_dims->dim(i)) +
+                      std::abs(stride_dims->dim(i)) - 1) /
+                     std::abs(stride_dims->dim(i));
+    // When begin tensor has negative values, currently range can't be computed.
+    if (begin_dims->dim(i) < 0) {
+      return errors::Unimplemented(
+          "Negative values in begin weight tensor are unsupported");
+    }
+    if (input_dims.dim_size(i) < 0) {
+      // end_dims and begin_dims do not have valid information yet.
+      dynamic_input_size_indices.push_back(i);
+    } else {
+      static_input_size_indices.push_back(i);
+      if (end_dims->dim(i) < begin_dims->dim(i) && stride_dims->dim(i) > 0) {
+        return errors::InvalidArgument(
+            "\"size\" cannot be negative for StridedSlice");
+      }
+    }
+  }
+
+  if (!dynamic_input_size_indices.empty()) {
+    if (strided_slice_spec == std::nullopt) {
+      return errors::InvalidArgument(
+          "The argument `strided_slice_spec` is "
+          "`std::nullopt` with `dynamic_input_size_indices` non empty.");
+    }
+    if (params->use_implicit_batch) {
+      return errors::InvalidArgument(
+          "In implicit batch mode, dynamic input size is not supported.");
+    }
+  }
+
+  if (params->validation_only) return OkStatus();
+
+  StatusOr<TRTNetworkBuilder> builder = TRTNetworkBuilder::Create(
+      params->converter->network(), params->weight_store);
+  TRT_ENSURE_OK(builder);
+
+  // Create the slice operation. For dynamic dims, the inputs of the operations
+  // may be reassigned later.
+  StatusOr<nvinfer1::ISliceLayer*> slice =
+      builder->Slice(input.tensor()->trt_tensor(), begin_dims->AsTrtDims(),
+                     size_dims, stride_dims->AsTrtDims());
+  TRT_ENSURE_PTR_OK(slice);
+
+  // Handle dynamic input shapes.
+  if (!dynamic_input_size_indices.empty()) {
+    TF_RETURN_IF_ERROR(HandleDynamicStridedSliceInput(
+        &*builder, *slice, *strided_slice_spec, dynamic_input_size_indices,
+        begin_dims->AsTrtDims(), stride_dims->AsTrtDims(),
+        end_dims->AsTrtDims()));
+  }
+
+  params->converter->SetLayerName(*slice, params->node_def, "slice",
+                                  op_instance);
+  ITensorProxyPtr tensor = (*slice)->getOutput(0);
+
+  // Reshape for shrink axis, ellipsis masks based on the shape computed by
+  // ValidateStridedSliceOp or HandleDynamicStridedSliceInput.
+  nvinfer1::Dims dims = tensor->trt_tensor()->getDimensions();
+  std::vector<int> slice_input_dims(dims.d, dims.d + dims.nbDims);
+  StridedSliceShapeSpec empty_spec;
+  empty_spec.shrink_axis_dense_mask = 0;
+  auto shrink_axis_mask =
+      strided_slice_spec.value_or(empty_spec).shrink_axis_dense_mask;
+  if (final_shape) {
+    if (shrink_axis_mask) {
+      int shrink_idx = params->use_implicit_batch ? 1 : 0;
+      const auto bShrink_axis_mask = std::bitset<32>(shrink_axis_mask);
+      for (int idx = 0; idx < slice_input_dims.size(); ++idx, ++shrink_idx) {
+        const bool shrink_axis = bShrink_axis_mask[shrink_idx];
+        if (shrink_axis) {
+          slice_input_dims[idx] = 0;
+        }
+      }
+      TF_RETURN_IF_ERROR(params->converter->SqueezeTensor(
+          tensor, &slice_input_dims, params, &tensor, op_instance));
+    } else {
+      /* To do: pmajety:
+            Remove the else condition when shrink_axis_mask is always defined */
+      TF_RETURN_IF_ERROR(PrepareTensorForShape(
+          params->converter, TRT_TensorOrWeights(tensor), *final_shape,
+          /*validation_only=*/false, &tensor, node_def, op_instance));
+    }
+  }
+  params->outputs->push_back(TRT_TensorOrWeights(tensor));
+  return OkStatus();
+}
+
+Status HandleDynamicStridedSliceInput(
+    TRTNetworkBuilder* builder, nvinfer1::ISliceLayer* slice_layer,
+    const StridedSliceShapeSpec& strided_slice_spec,
+    const absl::InlinedVector<int64, 4>& dynamic_input_size_indices,
+    nvinfer1::Dims begin_dims, nvinfer1::Dims stride_dims,
+    nvinfer1::Dims end_dims) {
+  TRT_ENSURE(builder);
+  TRT_ENSURE(slice_layer);
+
+  nvinfer1::ITensor* input_tensor = slice_layer->getInput(0);
+  TRT_ENSURE(input_tensor);
+
+  // When begin_mask or end_mask are set, we have to disregard the begin_tensor
+  // and end_tensor values. In static indices cases, ValidateStridedSliceOp
+  // returns the correct begin_tensor and end_tensor values, however with
+  // dynamic indices the correct shape has to be computed.
+
+  VLOG(3) << "begin_dims before: " << DebugString(begin_dims);
+  VLOG(3) << "end_dims before: " << DebugString(end_dims);
+  const auto begin_mask = std::bitset<32>(strided_slice_spec.begin_dense_mask);
+  const auto end_mask = std::bitset<32>(strided_slice_spec.end_dense_mask);
+  const auto shrink_axis_mask =
+      std::bitset<32>(strided_slice_spec.shrink_axis_dense_mask);
+  nvinfer1::Dims dims = input_tensor->getDimensions();
+
+  for (int idx = 0; idx < dims.nbDims; ++idx) {
+    VLOG(3) << "begin_mask[" << idx << "]: " << begin_mask[idx];
+    VLOG(3) << "end_mask[" << idx << "]: " << end_mask[idx];
+    VLOG(3) << "shrink_mask[" << idx << "]: " << shrink_axis_mask[idx];
+    if (begin_mask[idx]) {
+      begin_dims.d[idx] = 0;
+    }
+    if (end_mask[idx]) {
+      end_dims.d[idx] = dims.d[idx];
+    }
+    if (shrink_axis_mask[idx]) {
+      end_dims.d[idx] = begin_dims.d[idx] + 1;
+    }
+  }
+
+  VLOG(2) << "begin_dims after shrink_axis_mask correction: "
+          << DebugString(begin_dims);
+  VLOG(2) << "end_dims after shrink_axis_mask correction: "
+          << DebugString(end_dims);
+
+  // For each dynamic input dimension of the input, do some preprocessing based
+  // on whether this dimension is set in "begin_mask" or "end_mask" and the sign
+  // of the dimension's stride value.
+  // When stride is negative:
+  //   - If "begin_mask[dynamic_idx]" is set, then we need to adjust the slice
+  //     start of dimension[i] to the dynamic size.
+  //   - If "end_mask[dynamic_idx]" is set, it suffices to set
+  //     end_dims[dynamic_idx] to -1.
+  // When stride is positive:
+  //   - If "begin_mask[dynamic_idx]" is set, it suffices to set
+  //     begin_dims[dynamic_idx] to zero.
+  //   - If "end_mask[dynamic_idx]" is set, we need to adjust slice end to the
+  //     dynamic size of dimension "dynamic_idx".
+  absl::InlinedVector<int64, 4> dynamic_begin_indices;
+  absl::InlinedVector<int64, 4> dynamic_end_indices;
+
+  for (int i = 0; i < dynamic_input_size_indices.size(); i++) {
+    auto dynamic_idx = dynamic_input_size_indices[i];
+    if (begin_mask[dynamic_idx]) {
+      begin_dims.d[dynamic_idx] = 0;
+      if (stride_dims.d[dynamic_idx] < 0) {
+        dynamic_begin_indices.push_back(dynamic_idx);
+      }
+    }
+    if (end_mask[dynamic_idx] && !shrink_axis_mask[dynamic_idx]) {
+      end_dims.d[dynamic_idx] = stride_dims.d[dynamic_idx] > 0 ? 0 : -1;
+      if (stride_dims.d[dynamic_idx] > 0) {
+        dynamic_end_indices.push_back(dynamic_idx);
+      }
+    }
+  }
+
+  VLOG(2) << " Dynamic begin indices: " << DebugString(dynamic_begin_indices)
+          << " Dynamic end indices: " << DebugString(dynamic_end_indices);
+
+  // Create ITensors for each of the begin/stride/end constants.
+  StatusOr<nvinfer1::IConstantLayer*> begin_const = builder->Constant(
+      std::vector<int>(begin_dims.d, begin_dims.d + begin_dims.nbDims));
+  TRT_ENSURE_PTR_OK(begin_const);
+  nvinfer1::ITensor* begin_tensor = (*begin_const)->getOutput(0);
+  StatusOr<nvinfer1::IConstantLayer*> stride_const = builder->Constant(
+      std::vector<int>(stride_dims.d, stride_dims.d + stride_dims.nbDims));
+  TRT_ENSURE_PTR_OK(stride_const);
+  StatusOr<nvinfer1::IConstantLayer*> end_const = builder->Constant(
+      std::vector<int>(end_dims.d, end_dims.d + end_dims.nbDims));
+  TRT_ENSURE_PTR_OK(end_const);
+  nvinfer1::ITensor* end_tensor = (*end_const)->getOutput(0);
+
+  // Make corrections based on the begin_mask/end_mask values.
+  if (dynamic_end_indices.size() > 0) {
+    StatusOr<nvinfer1::IGatherLayer*> dynamic_end_masked_tensor =
+        builder->GetPartialShapeOf(input_tensor, dynamic_end_indices,
+                                   /*sub_one=*/false);
+    TRT_ENSURE_PTR_OK(dynamic_end_masked_tensor);
+    StatusOr<nvinfer1::IElementWiseLayer*> end_corrected =
+        builder->Add((*dynamic_end_masked_tensor)->getOutput(0), end_tensor);
+    TRT_ENSURE_PTR_OK(end_corrected);
+    end_tensor = (*end_corrected)->getOutput(0);
+  }
+  if (dynamic_begin_indices.size() > 0) {
+    StatusOr<nvinfer1::IGatherLayer*> dynamic_begin_masked_tensor =
+        builder->GetPartialShapeOf(input_tensor, dynamic_begin_indices,
+                                   /*sub_one=*/true);
+    TRT_ENSURE_PTR_OK(dynamic_begin_masked_tensor);
+
+    // Add back the original "begin" values for static dimensions.
+    StatusOr<nvinfer1::IElementWiseLayer*> begin_corrected = builder->Add(
+        (*dynamic_begin_masked_tensor)->getOutput(0), begin_tensor);
+    TRT_ENSURE_PTR_OK(begin_corrected);
+    begin_tensor = (*begin_corrected)->getOutput(0);
+  }
+
+  // Calculate the final size of the slice dynamically.
+  nvinfer1::ITensor* size_tensor;
+  {
+    StatusOr<nvinfer1::IElementWiseLayer*> num =
+        builder->Sub(end_tensor, begin_tensor);
+    TRT_ENSURE_PTR_OK(num);
+    StatusOr<nvinfer1::IElementWiseLayer*> ceil_div = builder->AbsCeilDivInt(
+        (*num)->getOutput(0), (*stride_const)->getOutput(0));
+    TRT_ENSURE_PTR_OK(ceil_div);
+    size_tensor = (*ceil_div)->getOutput(0);
+  }
+
+  slice_layer->setInput(1, *begin_tensor);
+  slice_layer->setInput(2, *size_tensor);
+  slice_layer->setInput(3, *(*stride_const)->getOutput(0));
+
+  return OkStatus();
+}
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/slice_ops.h

@@ -0,0 +1,70 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_OPS_SLICE_OPS_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_OPS_SLICE_OPS_H_
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/util/strided_slice_op.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+using SliceDims = absl::InlinedVector<int64, 4>;
+
+// Creates a strided slice operation using the given information. This function
+// expects that the begin, stride, and end vectors have already been validated.
+// This function converts the [begin:stride:end] specification to the TensorRT
+// [begin:stride:size] ISliceLayer specification. The following algorithm is
+// used to perform this conversion: 1) The given (input_dims,
+// [begin:stride:end]) specification is dividied into
+//  "static dimensions" and "dynamic dimensions". "Dynamic dimensions"
+//  includes all dimensions of the slice where input_dims[i] == -1.
+// 2a) If there are no dynamic dimensions, then the "begin", "stride", and
+//  "size" variables are passed to the ISLiceLayer creation as build-time
+//  constants in the form of nvinfer1::Dims objects.
+// 2b) If there are any dynamic dimensions, then the "begin", "stride", and
+//  "size" variables are treated as runtime dynamic shape Tensors in the
+//  TensorRT graph. In this case, we must calculate "size" at runtime for all
+//  dynamic dimensions, while static dimensions use the constant values.
+//
+// Note that when any dynamic indices are present (2b), the "strided_slice_spec"
+// must be specified. This structure can be obtained through the
+// "tensorflow::ValidateStridedSliceOp" function, or it can be constructed
+// directly. When the ValidateStridedSliceOp helper function is used, it will
+// also return the "begin", "stride", and "end" vectors. When all dimensions are
+// static (2a), the "strided_slice_spec" variable is not required.
+//
+// If the "final_shape" variable is specified, then a reshape operation will be
+// added to the graph to achieve this shape. The shape must be fully specified.
+//
+// "op_instance" is only required if the caller needs to pass this variable
+// through to the Converter functions optionally accept it (SetLayerName,
+// PrepareTensorForShape).
+Status ConvertStridedSliceHelper(
+    const OpConverterParams* params, const TRT_TensorOrWeights& input,
+    const PartialTensorShape& input_dims, const SliceDims& begin,
+    const SliceDims& stride, const SliceDims& end,
+    std::optional<nvinfer1::Dims> final_shape = std::nullopt,
+    std::optional<int> op_instance = std::nullopt,
+    std::optional<StridedSliceShapeSpec> strided_slice_spec = std::nullopt);
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_OPS_SLICE_OPS_H_

tensorflow/compiler/tf2tensorrt/convert/ops/softmax.cc

@@ -0,0 +1,81 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+class ConvertSoftmax : public OpConverterBase<ConvertSoftmax> {
+ public:
+  explicit ConvertSoftmax(const OpConverterParams *params)
+      : OpConverterBase<ConvertSoftmax>(params) {}
+
+  static constexpr std::array<DataType, 3> AllowedDataTypes() {
+    return {DataType::DT_FLOAT, DataType::DT_HALF};
+  }
+
+  static constexpr std::array<InputArgSpec, 1> InputSpec() {
+    return std::array<InputArgSpec, 1>{
+        InputArgSpec::Create("logits", TrtInputArg::kTensor)};
+  }
+
+  Status Validate() {
+    const auto &params = *this->params_;
+    const auto &inputs = params.inputs;
+
+    ITensorProxyPtr logits_tensor = inputs.at(0).tensor();
+    const int num_trt_dims = logits_tensor->getDimensions().nbDims;
+    if (!num_trt_dims && params.use_implicit_batch) {
+      return errors::InvalidArgument(
+          "TensorRT Softmax cannot apply on the batch dimension");
+    }
+    return OkStatus();
+  }
+
+  Status Convert() {
+    const auto &params = *this->params_;
+    const auto &inputs = params.inputs;
+    const auto &node_def = params.node_def;
+
+    ITensorProxyPtr logits_tensor = inputs.at(0).tensor();
+    const int num_trt_dims = logits_tensor->getDimensions().nbDims;
+
+    // Perform Softmax operation:
+    nvinfer1::ISoftMaxLayer *layer =
+        params.converter->network()->addSoftMax(*logits_tensor->trt_tensor());
+    TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name());
+    params.converter->SetLayerName(layer, node_def);
+    // Tensorflow SoftMax applies softmax operation over the last dimension.
+    layer->setAxes(1 << (num_trt_dims - 1));
+
+    ITensorProxyPtr output_tensor = layer->getOutput(0);
+    params.outputs->push_back(TRT_TensorOrWeights(output_tensor));
+    return OkStatus();
+  }
+};
+
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertSoftmax>(),
+                                  "Softmax");
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/tile.cc

@@ -0,0 +1,207 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+class ConvertTile : public OpConverterBase<ConvertTile> {
+ public:
+  explicit ConvertTile(const OpConverterParams *params)
+      : OpConverterBase<ConvertTile>(
+            params,
+            {DataType::DT_FLOAT, DataType::DT_HALF, DataType::DT_INT32}) {}
+
+  static constexpr std::array<InputArgSpec, 2> InputSpec() {
+    return std::array<InputArgSpec, 2>{
+        InputArgSpec::Create("input_tensor", TrtInputArg::kBoth),
+        InputArgSpec::Create("weight", TrtInputArg::kBoth)};
+  }
+
+  Status Validate() {
+    const auto &params = *this->params_;
+    const auto &inputs = params.inputs;
+
+    const auto &repl = inputs.at(1);
+    if (params.use_implicit_batch && repl.is_tensor()) {
+      return errors::InvalidArgument(
+          "Conversion for Tile is not implemented for multipliers "
+          "passed as a tensor in implicit batch mode.");
+    }
+
+    nvinfer1::DataType dtype;
+    const int *multiplies;
+    if (repl.is_weights()) {
+      TFTRT_CHECK_SHAPE_TENSOR(repl.weights().GetTensor());
+      dtype = repl.weights().TrtDType();
+      multiplies = repl.weights().GetPointer<int>();
+    } else {
+      dtype = repl.tensor()->getType();
+      multiplies = nullptr;
+    }
+
+    const auto &node = params.node_def;
+    TF_RETURN_IF_ERROR(check_type(dtype, nvinfer1::DataType::kINT32, node, 1));
+
+    const auto dims = inputs.at(0).GetTrtDims();
+    const auto nb_dims =
+        dims.nbDims +
+        (params.use_implicit_batch && inputs.at(0).is_tensor() ? 1 : 0);
+    if (multiplies) {
+      const int mult_numb = repl.weights().count();
+      if (mult_numb != nb_dims) {
+        return errors::InvalidArgument(
+            "The length of the replication vector (", mult_numb,
+            ") of the Tile operation in '", node.name(),
+            "' is expected to be equal to the rank of the input vector (",
+            nb_dims, ").");
+      }
+
+      if (std::any_of(multiplies, multiplies + nb_dims,
+                      [](int i) { return i <= 0; })) {
+        const auto &mul = absl::StrJoin(multiplies, multiplies + nb_dims, ", ");
+        return errors::InvalidArgument(
+            "All replications of the Tile operation in '", node.name(),
+            "' should be positive, got (", mul, ").");
+      }
+
+      if (params.use_implicit_batch && multiplies[0] > 1) {
+        return errors::Unimplemented(
+            "The Tile operation along the batch dimension in '", node.name(),
+            "' is not implemented.");
+      }
+    } else {
+      const auto &repl_dims = repl.GetTrtDims();
+      if (repl_dims.nbDims != 1) {
+        return errors::InvalidArgument(
+            "When replications are defined as a tensor, that tensor must be "
+            "1-dimensional. Got ",
+            repl_dims.nbDims, "-dimensional tensor.");
+      }
+
+      // Check the number of elements in multiplyer for tensors with non-dynamic
+      // shape
+      if (repl_dims.d[0] >= 0 && repl_dims.d[0] != nb_dims) {
+        return errors::InvalidArgument(
+            "When replications are defined as a tensor, "
+            "the number of its elements (",
+            repl_dims.d[0], ") must be equal to the rank of the input tensor (",
+            nb_dims, ").");
+      }
+    }
+
+    return OkStatus();
+  }
+
+  Status Convert() {
+    const auto &params = *this->params_;
+    const auto &inputs = params.inputs;
+    auto *converter = params.converter;
+    auto *network = converter->network();
+    const auto &tensor = inputs.at(0);
+    const auto &replics = inputs.at(1);
+    const auto dims = tensor.GetTrtDims();
+    const auto nb_dims = dims.nbDims;
+
+    nvinfer1::Dims output_size{nb_dims, {1}};
+    bool dynamic_flag = replics.is_tensor() || !HasStaticShape(dims);
+
+    if (!dynamic_flag) {
+      // If input0 is a tensor, and we're in implicit batch mode, then we need
+      // dim_offset.
+      const auto dim_offset =
+          params.use_implicit_batch && tensor.is_tensor() ? 1 : 0;
+      const auto *input_size = dims.d;
+      const int *pReplics = replics.weights().GetPointer<int>() + dim_offset;
+      for (int i = 0; i < nb_dims; i++)
+        output_size.d[i] = pReplics[i] * input_size[i];
+    }
+
+    StatusOr<TRTNetworkBuilder> builder;
+    if (tensor.is_weights() || (dynamic_flag && replics.is_weights())) {
+      builder =
+          TRTNetworkBuilder::Create(converter->network(), params.weight_store);
+      TRT_ENSURE_OK(builder);
+    }
+
+    ITensorProxyPtr input_tensor;
+    if (tensor.is_weights()) {
+      StatusOr<nvinfer1::IConstantLayer *> weights_const =
+          builder->WeightsToConstant(tensor.weights().GetTrtWeights(), dims);
+      TRT_ENSURE_PTR_OK(weights_const);
+      input_tensor = (*weights_const)->getOutput(0);
+    } else {
+      input_tensor = tensor.tensor();
+    }
+
+    auto &input_trt_tensor = *input_tensor->trt_tensor();
+    nvinfer1::ITensor *target_shape = nullptr;
+    if (dynamic_flag) {
+      nvinfer1::ITensor *mult;
+      if (replics.is_weights()) {
+        StatusOr<nvinfer1::IConstantLayer *> weights_const =
+            builder->WeightsToConstant(replics.weights().GetTrtWeights(),
+                                       replics.GetTrtDims());
+        TRT_ENSURE_PTR_OK(weights_const);
+        mult = (*weights_const)->getOutput(0);
+      } else {
+        const ITensorProxyPtr multiplies = replics.tensor()->trt_tensor();
+        mult = multiplies->trt_tensor();
+      }
+
+      nvinfer1::ITensor *shape =
+          network->addShape(input_trt_tensor)->getOutput(0);
+      target_shape = network
+                         ->addElementWise(*shape, *mult,
+                                          nvinfer1::ElementWiseOperation::kPROD)
+                         ->getOutput(0);
+    }
+
+    nvinfer1::Dims start{nb_dims, {}};
+    DimsAdapter stride(std::vector<int>(nb_dims, 1));
+    auto layer = network->addSlice(input_trt_tensor, start, output_size,
+                                   stride.AsTrtDims());
+    layer->setMode(nvinfer1::SliceMode::kWRAP);
+    if (target_shape) layer->setInput(2, *target_shape);
+
+    converter->SetLayerName(layer, params.node_def.name(), "to_tile");
+    ITensorProxyPtr output_tensor = layer->getOutput(0);
+    if (tensor.is_weights() && params.use_implicit_batch) {
+      // Reshape output tensor by removing first dimension.
+      DimsAdapter adap(output_tensor->getDimensions());
+      TF_RETURN_IF_ERROR(adap.RemoveBatchDimension());
+
+      TF_RETURN_IF_ERROR(PrepareTensorForShape(
+          params.converter, TRT_TensorOrWeights(output_tensor),
+          adap.AsTrtDims(), false, &output_tensor, params.node_def));
+    }
+
+    AddOutput(TRT_TensorOrWeights(output_tensor));
+    return OkStatus();
+  }
+};
+
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertTile>(), "Tile");
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/unary_ops.cc

@@ -0,0 +1,251 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+const UnaryOperationMapType* UnaryOperationMap() {
+  static auto* const m = new UnaryOperationMapType({
+    {"Exp", nvinfer1::UnaryOperation::kEXP},
+        {"Log", nvinfer1::UnaryOperation::kLOG},
+        {"Sqrt", nvinfer1::UnaryOperation::kSQRT},
+        {"Rsqrt", nvinfer1::UnaryOperation::kSQRT},
+        {"Reciprocal", nvinfer1::UnaryOperation::kRECIP},
+        {"Abs", nvinfer1::UnaryOperation::kABS},
+        {"Neg", nvinfer1::UnaryOperation::kNEG},
+        {"Sin", nvinfer1::UnaryOperation::kSIN},
+        {"Cos", nvinfer1::UnaryOperation::kCOS},
+        {"Tan", nvinfer1::UnaryOperation::kTAN},
+        {"Sinh", nvinfer1::UnaryOperation::kSINH},
+        {"Cosh", nvinfer1::UnaryOperation::kCOSH},
+        {"Asin", nvinfer1::UnaryOperation::kASIN},
+        {"Acos", nvinfer1::UnaryOperation::kACOS},
+        {"Atan", nvinfer1::UnaryOperation::kATAN},
+        {"Asinh", nvinfer1::UnaryOperation::kASINH},
+        {"Acosh", nvinfer1::UnaryOperation::kACOSH},
+        {"Atanh", nvinfer1::UnaryOperation::kATANH},
+        {"Ceil", nvinfer1::UnaryOperation::kCEIL},
+        {"Floor", nvinfer1::UnaryOperation::kFLOOR},
+        {"Erf", nvinfer1::UnaryOperation::kERF},
+#if IS_TRT_VERSION_GE(8, 2, 0, 0)
+        {"Round", nvinfer1::UnaryOperation::kROUND},
+        {"Sign", nvinfer1::UnaryOperation::kSIGN},
+#endif
+  });
+  return m;
+}
+
+const UnaryOperationMapType* UnaryBooleanOperationMap() {
+  static auto* const m = new UnaryOperationMapType({
+      {"LogicalNot", nvinfer1::UnaryOperation::kNOT},
+  });
+  return m;
+}
+
+const ActivationTypeMapType* ActivationTypeMap() {
+  static auto* const m = new ActivationTypeMapType({
+      {"LeakyRelu", nvinfer1::ActivationType::kLEAKY_RELU},
+      {"Relu", nvinfer1::ActivationType::kRELU},
+      {"Relu6", nvinfer1::ActivationType::kCLIP},
+      {"Sigmoid", nvinfer1::ActivationType::kSIGMOID},
+      {"Tanh", nvinfer1::ActivationType::kTANH},
+      {"Elu", nvinfer1::ActivationType::kELU},
+      {"Selu", nvinfer1::ActivationType::kSELU},
+      {"Softsign", nvinfer1::ActivationType::kSOFTSIGN},
+      {"Softplus", nvinfer1::ActivationType::kSOFTPLUS},
+  });
+  return m;
+}
+
+template <typename T>
+class ConvertUnaryImpl {
+ protected:
+  ConvertUnaryImpl(const OperationMap<T>* pOperMap) : pOperMap_(pOperMap) {}
+
+  Status ValidateImpl(const OpConverterParams& params,
+                      const std::vector<string>& not_supported_ops = {}) {
+    const auto& node = params.node_def;
+    const auto& op = node.op();
+    if (pOperMap_->find(op) == pOperMap_->end()) {
+      return errors::Unimplemented("Unary op: ", op, " not supported");
+    }
+    DimsAdapter input_dims(params.inputs.at(0).GetTrtDims());
+    if (!input_dims.NumDims()) {
+      return errors::InvalidArgument(
+          "At least 1 dimension is required for UNARY operation '", op, "'");
+    }
+
+    if (!not_supported_ops.empty() && params.use_implicit_batch) {
+      const auto& end = not_supported_ops.end();
+      if (std::find(not_supported_ops.begin(), end, op) != end) {
+        const auto& err =
+            convert_not_supported_implicit(op, node.name(), "Unary");
+        return errors::Unimplemented(err);
+      }
+    }
+
+    return OkStatus();
+  }
+
+  Status ConvertImpl(const OpConverterParams& params) {
+    const auto& node_def = params.node_def;
+    auto* converter = params.converter;
+    const auto op_pair = pOperMap_->find(node_def.op());
+    ITensorProxyPtr tensor = params.inputs.at(0).tensor();
+    nvinfer1::IUnaryLayer* layer =
+        converter->network()->addUnary(*tensor->trt_tensor(), op_pair->second);
+    TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name());
+    converter->SetLayerName(layer, node_def);
+    if (node_def.op() == "Rsqrt") {
+      layer = converter->network()->addUnary(*layer->getOutput(0),
+                                             nvinfer1::UnaryOperation::kRECIP);
+      TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name());
+      converter->SetLayerName(layer, node_def, "recip");
+    }
+    params.outputs->push_back(TRT_TensorOrWeights(layer->getOutput(0)));
+    return OkStatus();
+  }
+  static constexpr std::array<InputArgSpec, 1> InputSpec() {
+    return std::array<InputArgSpec, 1>{
+        InputArgSpec::Create("x", TrtInputArg::kTensor)};
+  }
+
+ protected:
+  const OperationMap<T>* pOperMap_;
+};
+
+class ConvertUnary : public OpConverterBase<ConvertUnary>,
+                     protected ConvertUnaryImpl<nvinfer1::UnaryOperation> {
+ public:
+  explicit ConvertUnary(const OpConverterParams* params)
+      : OpConverterBase<ConvertUnary>(
+            params,
+            params->node_def.op() == "Sign"
+                ? std::vector<DataType>{DataType::DT_FLOAT, DataType::DT_HALF,
+                                        DataType::DT_INT8, DT_INT32}
+                : std::vector<DataType>{DataType::DT_FLOAT, DataType::DT_HALF,
+                                        DataType::DT_INT8}),
+        ConvertUnaryImpl(UnaryOperationMap()) {}
+
+  static constexpr std::array<InputArgSpec, 1> InputSpec() {
+    return ConvertUnaryImpl::InputSpec();
+  }
+
+  Status Validate() { return ValidateImpl(*params_, {"Sign", "Round"}); }
+  Status Convert() { return ConvertImpl(*params_); }
+};
+
+class ConvertBooleanUnary : public OpConverterBase<ConvertBooleanUnary>,
+                            public ConvertUnaryImpl<nvinfer1::UnaryOperation> {
+ public:
+  explicit ConvertBooleanUnary(const OpConverterParams* params)
+      : OpConverterBase<ConvertBooleanUnary>(params, {DataType::DT_BOOL}),
+        ConvertUnaryImpl(UnaryBooleanOperationMap()) {}
+
+  static constexpr std::array<InputArgSpec, 1> InputSpec() {
+    return ConvertUnaryImpl::InputSpec();
+  }
+
+  static constexpr const char* NodeDefDataTypeAttributeName() {
+    /*
+    node {
+      name: "..."
+      op: "LogicalNot"
+      input: "..."
+    }
+    */
+    return "";
+  }
+  Status Validate() {
+#if IS_TRT_VERSION_GE(8, 2, 0, 0)
+    return ValidateImpl(*params_, {"LogicalNot"});
+#else
+    return errors::Unimplemented("Boolean op: ", params_->node_def.op(),
+                                 " is not supported in TRT version < 8.2");
+#endif
+  }
+  Status Convert() { return ConvertImpl(*params_); }
+};
+
+class ConvertActivation : public OpConverterBase<ConvertActivation>,
+                          protected ConvertUnaryImpl<nvinfer1::ActivationType> {
+ public:
+  explicit ConvertActivation(const OpConverterParams* params)
+      : OpConverterBase<ConvertActivation>(params),
+        ConvertUnaryImpl(ActivationTypeMap()) {}
+
+  static constexpr std::array<InputArgSpec, 1> InputSpec() {
+    return std::array<InputArgSpec, 1>{
+        InputArgSpec::Create("input", TrtInputArg::kTensor)};
+  }
+
+  Status Validate() {
+    TF_RETURN_IF_ERROR(ValidateImpl(*params_));
+    const auto& node_def = params_->node_def;
+    if (node_def.op() == "LeakyRelu") {
+      return GetNodeAttr(AttrSlice(node_def), "alpha", &alpha_);
+    }
+    alpha_ = 1.0f;
+    return OkStatus();
+  }
+  Status Convert() {
+    auto* converter = params_->converter;
+    const auto& inputs = params_->inputs;
+    const auto& node_def = params_->node_def;
+    const auto& op = node_def.op();
+    const auto op_pair = pOperMap_->find(op);
+    nvinfer1::IActivationLayer* layer = converter->network()->addActivation(
+        *inputs.at(0).tensor()->trt_tensor(), op_pair->second);
+    TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name());
+    converter->SetLayerName(layer, node_def, "activation");
+    ITensorProxyPtr output_tensor = layer->getOutput(0);
+    // Set parameters.
+    if (op == "Selu") {
+      // From tensorflow/core/kernels/relu_op_functor.h
+      alpha_ = 1.7580993408473768599402175208123f;
+      layer->setBeta(1.0507009873554804934193349852946f);
+    } else if (op == "Softplus") {
+      layer->setBeta(1.0f);
+    } else if (op == "Relu6") {
+      layer->setBeta(6.0f);
+      converter->ProvideQuantizationRange(&output_tensor, alpha_ = 0.0f, 6.0f);
+    }
+    layer->setAlpha(alpha_);
+    params_->outputs->push_back(TRT_TensorOrWeights(output_tensor));
+    return OkStatus();
+  }
+
+ private:
+  float alpha_ = 0.f;
+};
+
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertUnary>(),
+                                  GetOperationNames(*UnaryOperationMap()));
+REGISTER_DEFAULT_TRT_OP_CONVERTER(
+    MakeConverterFunction<ConvertBooleanUnary>(),
+    GetOperationNames(*UnaryBooleanOperationMap()));
+
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertActivation>(),
+                                  GetOperationNames(*ActivationTypeMap()));
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

tensorflow/compiler/tf2tensorrt/convert/ops/variable_ops.cc

@@ -0,0 +1,356 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
+#include "tensorflow/core/common_runtime/process_function_library_runtime.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/stream_executor.h"
+#include "third_party/tensorrt/NvInfer.h"
+#include "third_party/tensorrt/NvInferRuntimeCommon.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+struct VarAttributes {
+  TensorShapeProto shape_proto;
+  TensorShape shape;
+  string name;
+  DataType dtype;
+  string shared_name;
+  string container;
+};
+
+template <typename T, bool is_resource>
+Status ReadVariableHelper(const OpConverterParams* params,
+                          const VarAttributes& attrs,
+                          TRT_ShapedWeights* weights) {
+  Tensor tensor(attrs.dtype, attrs.shape);
+  auto ctx = params->converter->context();
+  TRT_ENSURE(ctx != nullptr);
+  auto tensor_flat = tensor.flat<T>();
+
+  // Clone function library runtime in order to get a mutable library
+  // definition to add and run a function with the variable operation.
+  auto lib = ctx->function_library();
+  std::unique_ptr<FunctionLibraryDefinition> lib_def;
+  std::unique_ptr<ProcessFunctionLibraryRuntime> lib_pflr;
+  FunctionLibraryRuntime* lib_clone;  // Not owned.
+  TF_RETURN_IF_ERROR(lib->Clone(&lib_def, &lib_pflr, &lib_clone));
+
+  // Create function definition.
+  FunctionDef fdef;
+  std::vector<Tensor> args;
+  string func_name = attrs.name + "/func";
+  if (is_resource) {
+    // Create input tensor with the resource handle.
+    const auto& inputs = params->inputs;
+    const TRT_TensorOrWeights& handle = inputs.at(0);
+    args.emplace_back(handle.resource());
+
+    fdef = FunctionDefHelper::Define(
+        func_name,                                             // Name
+        {"in: resource"},                                      // Args
+        {absl::StrCat("out: ", DataTypeString(attrs.dtype))},  // Returns
+        {},                                                    // Attr def
+        // Nodes
+        {{{attrs.name},
+          "ReadVariableOp",
+          {"in"},  // Name of the Placeholder or VarHandleOp
+          {{"dtype", attrs.dtype}}},
+         {{"out"}, "Identity", {attrs.name}, {{"T", attrs.dtype}}}});
+  } else {
+    fdef = FunctionDefHelper::Define(
+        func_name,                                             // Name
+        {},                                                    // Args
+        {absl::StrCat("out: ", DataTypeString(attrs.dtype))},  // Returns
+        {},                                                    // Attr def
+        // Nodes
+        {{{attrs.name},
+          "VariableV2",
+          {},
+          {{"dtype", attrs.dtype},
+           {"shape", attrs.shape_proto},
+           {"container", attrs.container},
+           {"shared_name", attrs.shared_name}}},
+         {{"out"}, "Identity", {attrs.name}, {{"T", attrs.dtype}}}});
+  }
+
+  // Add function definition to the library.
+  TF_RETURN_IF_ERROR(lib_def->AddFunctionDef(fdef));
+
+  // Instantiate function.
+  FunctionLibraryRuntime::Handle func_handle;
+  FunctionLibraryRuntime::InstantiateOptions inst_ops;
+  inst_ops.state_handle = "";
+  inst_ops.target = ctx->device()->name();
+  AttrValueMap attr_list;
+  TF_RETURN_IF_ERROR(lib_clone->Instantiate(func_name, AttrSlice(&attr_list),
+                                            inst_ops, &func_handle));
+
+  FunctionLibraryRuntime::Options opts;
+  opts.rendezvous = ctx->rendezvous();
+  opts.cancellation_manager = ctx->cancellation_manager();
+  opts.runner = ctx->runner();
+
+  std::vector<Tensor>* rets = new std::vector<Tensor>();
+  std::unique_ptr<std::vector<Tensor>> outputs_wrapper(rets);
+
+  // Run the new function synchronously.
+  TF_RETURN_IF_ERROR(lib_clone->RunSync(opts, func_handle, args, rets));
+
+  TRT_ENSURE(ctx->op_device_context() != nullptr);
+  TRT_ENSURE(ctx->op_device_context()->stream() != nullptr);
+
+  // Copy tensor.
+  cudaStream_t stream = reinterpret_cast<cudaStream_t>(CHECK_NOTNULL(
+      ctx->op_device_context()->stream()->platform_specific_handle().stream));
+
+  auto ret = cudaMemcpyAsync(tensor_flat.data(), rets->at(0).flat<T>().data(),
+                             rets->at(0).NumElements() * sizeof(T),
+                             cudaMemcpyDeviceToHost, stream);
+  if (ret != 0) {
+    return errors::Internal("Could not copy the variable ", attrs.name);
+  }
+  cudaStreamSynchronize(stream);
+
+  TF_RETURN_IF_ERROR(
+      TfTensorToTrtWeights(tensor, params->weight_store, weights));
+
+  return OkStatus();
+}
+
+class ConvertVariableV2 : public OpConverterBase<ConvertVariableV2> {
+ public:
+  ConvertVariableV2(const OpConverterParams* params)
+      : OpConverterBase<ConvertVariableV2>(params) {}
+
+  static constexpr std::array<InputArgSpec, 0> InputSpec() { return {}; }
+
+  static constexpr const char* NodeDefDataTypeAttributeName() {
+    /*
+    node {
+      name: "..."
+      op: "VariableV2"
+      ...
+      attr {
+        key: "dtype"
+        value {
+          type: DT_FLOAT
+        }
+      }
+      ...
+    }
+    */
+    return "dtype";
+  }
+
+  template <typename T>
+  Status ValidateImpl() {
+    const auto& node_def = params_->node_def;
+
+    // Verify and consume node attributes.
+    StatusOr<TensorShapeProto> shape_proto =
+        GetAttrValue<TensorShapeProto>("shape");
+    StatusOr<string> shared_name = GetAttrValue<string>("shared_name");
+    StatusOr<string> container = GetAttrValue<string>("container");
+    TRT_ENSURE_OK(shape_proto);
+    TRT_ENSURE_OK(shared_name);
+    TRT_ENSURE_OK(container);
+
+    attrs_.shape_proto = *shape_proto;
+    attrs_.shape = TensorShape(*shape_proto);
+    attrs_.name = node_def.name();
+    attrs_.shared_name = *shared_name;
+    attrs_.container = *container;
+
+    Tensor tensor(attrs_.dtype, attrs_.shape);
+    auto tensor_flat = tensor.flat<T>();
+    for (int64_t i = 0; i < tensor_flat.size(); i++) {
+      tensor_flat(i) = T(0.0f);
+    }
+
+    TRT_ShapedWeights weights;
+    TF_RETURN_IF_ERROR(
+        TfTensorToTrtWeights(tensor, params_->weight_store, &weights));
+
+    // Only push outputs during validation and when outputs are expected.
+    if (params_->validation_only && params_->outputs != nullptr) {
+      AddOutput(TRT_TensorOrWeights(weights));
+    }
+    return OkStatus();
+  }
+
+  Status Validate() {
+    const auto& node_def = params_->node_def;
+    StatusOr<DataType> dtype = GetAttrValue<DataType>("dtype");
+    TRT_ENSURE_OK(dtype);
+    attrs_.dtype = *dtype;
+
+    switch (attrs_.dtype) {
+      case DT_FLOAT:
+        return ValidateImpl<float>();
+      case DT_HALF:
+        return ValidateImpl<Eigen::half>();
+      default:
+        // Note: this should have been caught by ValidateNodeDefDataType, but
+        // the compiler expects that all paths be handled in switch.
+        return errors::Unimplemented("Data type ", DataTypeString(attrs_.dtype),
+                                     " is not supported for ", node_def.op(),
+                                     ", at ", node_def.name());
+    }
+  }
+
+  template <typename T>
+  Status ConvertImpl() {
+    TRT_ShapedWeights weights;
+    TF_RETURN_IF_ERROR(ReadVariableHelper<T, false>(params_, attrs_, &weights));
+    AddOutput(TRT_TensorOrWeights(weights));
+    return OkStatus();
+  }
+
+  Status Convert() {
+    const auto& node_def = params_->node_def;
+
+    switch (attrs_.dtype) {
+      case DT_FLOAT:
+        return ConvertImpl<float>();
+      case DT_HALF:
+        return ConvertImpl<Eigen::half>();
+      default:
+        // Note: this should have been caught by ValidateNodeDefDataType, but
+        // the compiler expects that all paths be handled in switch.
+        return errors::Unimplemented("Data type ", DataTypeString(attrs_.dtype),
+                                     " is not supported for ", node_def.op(),
+                                     ", at ", node_def.name());
+    }
+  }
+
+ private:
+  VarAttributes attrs_{};
+};
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertVariableV2>(),
+                                  {"VariableV2"});
+
+class ConvertReadVariableOp : public OpConverterBase<ConvertReadVariableOp> {
+ public:
+  ConvertReadVariableOp(const OpConverterParams* params)
+      : OpConverterBase<ConvertReadVariableOp>(params) {}
+
+  static constexpr std::array<InputArgSpec, 1> InputSpec() {
+    return {InputArgSpec::Create("resource", TrtInputArg::kResource)};
+  }
+
+  static constexpr const char* NodeDefDataTypeAttributeName() {
+    return "dtype";
+  }
+
+  template <typename T>
+  Status ValidateImpl() {
+    const auto& node_def = params_->node_def;
+
+    // Verify and consume node attributes.
+    StatusOr<TensorShapeProto> shape_proto =
+        GetAttrValue<TensorShapeProto>("_shape");
+    TRT_ENSURE_OK(shape_proto);
+
+    attrs_.shape_proto = *shape_proto;
+    attrs_.shape = TensorShape(*shape_proto);
+    attrs_.name = node_def.name();
+
+    Tensor tensor(attrs_.dtype, attrs_.shape);
+    auto tensor_flat = tensor.flat<T>();
+    for (int64_t i = 0; i < tensor_flat.size(); i++) {
+      tensor_flat(i) = T(0.0f);
+    }
+
+    TRT_ShapedWeights weights;
+    TF_RETURN_IF_ERROR(
+        TfTensorToTrtWeights(tensor, params_->weight_store, &weights));
+
+    // Only push outputs during validation and when outputs are expected.
+    if (params_->validation_only && params_->outputs != nullptr) {
+      AddOutput(TRT_TensorOrWeights(weights));
+    }
+    return OkStatus();
+  }
+
+  Status Validate() {
+    const auto& node_def = params_->node_def;
+    if (params_->use_implicit_batch) {
+      return errors::Unimplemented("Implicit batch mode not supported, at ",
+                                   node_def.name());
+    }
+
+    StatusOr<DataType> dtype = GetAttrValue<DataType>("dtype");
+    TRT_ENSURE_OK(dtype);
+    attrs_.dtype = *dtype;
+
+    switch (attrs_.dtype) {
+      case DT_FLOAT:
+        return ValidateImpl<float>();
+      case DT_HALF:
+        return ValidateImpl<Eigen::half>();
+      default:
+        // Note: this should have been caught by ValidateNodeDefDataType, but
+        // the compiler expects that all paths be handled in switch.
+        return errors::Unimplemented("Data type ", DataTypeString(attrs_.dtype),
+                                     " is not supported for ", node_def.op(),
+                                     ", at ", node_def.name());
+    }
+  }
+
+  template <typename T>
+  Status ConvertImpl() {
+    TRT_ShapedWeights weights;
+    TF_RETURN_IF_ERROR(ReadVariableHelper<T, true>(params_, attrs_, &weights));
+    AddOutput(TRT_TensorOrWeights(weights));
+    return OkStatus();
+  }
+
+  Status Convert() {
+    const auto& node_def = params_->node_def;
+
+    switch (attrs_.dtype) {
+      case DT_FLOAT:
+        return ConvertImpl<float>();
+      case DT_HALF:
+        return ConvertImpl<Eigen::half>();
+      default:
+        // Note: this should have been caught by ValidateNodeDefDataType, but
+        // the compiler expects that all paths be handled in switch.
+        return errors::Unimplemented("Data type ", DataTypeString(attrs_.dtype),
+                                     " is not supported for ", node_def.op(),
+                                     ", at ", node_def.name());
+    }
+  }
+
+ private:
+  VarAttributes attrs_{};
+};
+REGISTER_DEFAULT_TRT_OP_CONVERTER(
+    MakeConverterFunction<ConvertReadVariableOp>(), {"ReadVariableOp"});
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT