/* 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 implict 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 implict 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_pod<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_pod<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* const network_;
  TrtWeightStore* const 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_
