Value output_layer_norm_coefficients_;

  mlir::TFL::LSTMOp lstm_;

  Value none_;
  SmallVector<int64_t, 1> bias_slice_shape_;
  SmallVector<int64_t, 1> bias_size_values_;
  SmallVector<int64_t, 2> weight_slice_shape_;
  SmallVector<int64_t, 2> weight_slice_size_input_values_;
  SmallVector<int64_t, 2> weight_slice_size_recurrent_values_;
  OpBuilder builder_;
};

      : input_shapes_(input_shapes) {}
  void runOnOperation() override {
    // Parse `input_arg_shapes_` if provided (test only)
    SmallVector<ArrayRef<int64_t>> input_shapes_vec;
    absl::StatusOr<SmallVector<SmallVector<int64_t>>> parsed_shapes;
    if (!input_arg_shapes_.empty()) {
      parsed_shapes = ParseArgumentShapes(input_arg_shapes_);
      if (!parsed_shapes.ok()) {

  // TFL::UnsortedSegmentProdOp.
  llvm::SmallVector<int32_t, 4> flattened_out_segids =
      llvm::SmallVector<int32_t, 4>(operand_rank, static_cast<int32_t>(-1));
  for (int64_t i : dot_dimensions_info.out_dimensions().AxesArray()) {
    flattened_out_segids[i] = 0;
  }
  llvm::SmallVector<int32_t, 4> flattened_contracting_segids =
      llvm::SmallVector<int32_t, 4>(operand_rank, static_cast<int32_t>(-1));

    StringAttr conv_padding, Value &padding, int num_dims) {
  Value zero_rank1 = CreateConstValue<int32_t>(builder, loc, {1}, {0});
  SmallVector<Value> temp_padding_values{zero_rank1, zero_rank1};

  auto reshape_op = [&](Value value, const SmallVector<int64_t> &shape) {
    const int64_t rank = shape.size();
    return builder.create<TF::ReshapeOp>(

    mlir::OpBuilder* builder,
    llvm::SmallVectorImpl<llvm::SmallVector<mlir::Value, 4>>* input_list);

// Extracts a list of OpSharding that represent output sharding configuration of
// `tf_device.cluster`.
mlir::LogicalResult ParseAndValidateOutputSharding(
    int num_cores_per_replica, mlir::tf_device::ClusterFuncOp cluster_func,
    mlir::SmallVector<xla::OpSharding, 4>* output_sharding_list);

// in 'arguments_to_erase'.
template <typename T, typename U>
T GetUpdatedWhileOp(T while_op, const U& argument_types,
                    const llvm::SmallVector<unsigned, 4>& arguments_to_erase) {
  OpBuilder builder(while_op);
  llvm::SmallVector<Type, 4> new_operand_types;
  llvm::SmallVector<Value> new_operands;
  auto operands = while_op->getOperands();
  const int num_operands = while_op->getNumOperands();

#include <type_traits>

#include "llvm/ADT/ArrayRef.h"  // IWYU pragma: keep; required to include the definition of ArrayRef
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"  // IWYU pragma: keep; required to include the definition of SmallVector
#include "mlir/Support/LLVM.h"  // from @llvm-project

namespace mlir::quant {

// Permutes `values` with `permutation`. Returns the permuted values. Sizes of

  SmallVector<float> TransposeValues(const ArrayRef<float> values) const {
    SmallVector<float> transposed_values(values.size());
    SmallVector<int64_t> current_indices = {};
    TransposeRecursively(values, transposed_values, current_indices);

    return transposed_values;
  }

  // Returns the shape after permutation.
  SmallVector<int64_t> GetTargetShape() const { return target_shape_; }

  } else {
    // Convert I64 const array to I32.
    llvm::SmallVector<int32_t> const_i32_vec;
    for (auto element : const_value) {
      const_i32_vec.push_back(static_cast<int32_t>(element));
    }
    const_value_raw = rewriter.getI32TensorAttr(const_i32_vec);
  }
  Value result_const = builder.create<TF::ConstOp>(const_value_raw);
  return result_const;
}

llvm::SmallVector<int64_t> GetInversePermutationArray(

void AssignDevicesToReplicate(
    tf_device::ReplicateOp replicate,
    llvm::ArrayRef<llvm::SmallVector<tensorflow::TPUDeviceAndHost, 8>>
        tpu_devices,
    OpBuilder* builder) {
  if (!replicate) return;

  const int num_replicas = tpu_devices.size();
  const int num_cores_per_replica = tpu_devices.front().size();

  llvm::SmallVector<NamedAttribute, 8> device_attrs;
  for (int core = 0; core < num_cores_per_replica; ++core) {