// inputs.
std::optional<llvm::StringRef> GetXlaShardingFromArg(
    Value value, const llvm::SmallVector<std::string>& logical_device_vec) {
  llvm::SmallPtrSet<Value, 4> visited_values;
  llvm::SmallVector<Value, 4> values_to_visit{value};
  while (!values_to_visit.empty()) {
    llvm::SmallVector<Value, 4> next_values_to_visit;
    for (Value value_to_visit : values_to_visit) {

  if (!first->hasOneUse() || !second->hasOneUse())
    return rewriter.notifyMatchFailure(concat, "slice ops are used elsewhere");

  SmallVector<int64_t> new_start;
  SmallVector<int64_t> new_limit;
  SmallVector<int64_t> new_slice_shape;
  new_start.reserve(first.getStrides().size());
  new_limit.reserve(first.getStrides().size());
  new_slice_shape.reserve(first.getStrides().size());

                                     zero_point, qmin, qmax);
  } else if (auto q_type = dyn_cast<quant::UniformQuantizedPerAxisType>(type)) {
    const int size = q_type.getScales().size();
    SmallVector<double, 4> scales(size);
    SmallVector<int64_t, 4> zero_points(size);
    for (int i = 0; i < size; ++i) {
      scales[i] = recalculate_scale(q_type.getScales()[i]);
      zero_points[i] = recalculate_zero_point(q_type.getZeroPoints()[i]);

  // Parse operands, attributes, and region of op.
  llvm::SmallVector<llvm::SmallVector<OpAsmParser::UnresolvedOperand, 8>, 8>
      replicated_inputs;
  llvm::SmallVector<OpAsmParser::UnresolvedOperand, 8> packed_inputs;
  llvm::SmallVector<OpAsmParser::UnresolvedOperand, 8> region_args;
  llvm::SmallVector<Type, 8> region_arg_types;
  int32_t n = 0;
  Region& body = *result.addRegion();

#include "tensorflow/compiler/mlir/tensorflow/transforms/collection_ops_util.h"

#include <optional>

#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/IR/Attributes.h"  // from @llvm-project
#include "mlir/IR/Builders.h"  // from @llvm-project

    // UnidirectionalSequenceLstm should only have 1 output, and that is the
    // original ophint converted node's 3rd output.
    SmallVector<Type, 4> result_types;
    result_types.push_back(op->getOpResult(2).getType());

    // Populate attributes.
    SmallVector<NamedAttribute, 4> attributes;
    // Activation will always be tanh.
    attributes.push_back(rewriter.getNamedAttr("fused_activation_function",

  std::vector<std::string> input_names, output_names;
  GetUniqueInputOutputNodeNames(module_op, input_names, output_names);

  // Collects argument and result types.
  llvm::SmallVector<Location> arg_locs;
  llvm::SmallVector<Type> arg_types, result_types;

  for (auto func_op : module_op.getOps<func::FuncOp>()) {
    if (!ShouldIncludeInMainFunction(func_op)) continue;

// tuple type, a new tuple type with `type` and `mhlo.token` type is created
// instead.
SmallVector<Type> GetTypeWithToken(OpBuilder& builder, ArrayRef<Type> types,
                                   bool flatten_tuple) {
  SmallVector<Type> new_result_types;
  auto token_type = TokenType::get(builder.getContext());

  if (flatten_tuple) {

  }

  return op->emitError() << "unknown tf_saved_model dialect attribute '"
                         << named_attr.getName().getValue() << "'";
}

SmallVector<StringRef, 2> GetExportedNames(Operation *op) {
  SmallVector<StringRef, 2> ret;
  auto exported_names =
      op->getAttrOfType<ArrayAttr>(kTfSavedModelExportedNamesAttr);
  if (exported_names) {
    for (auto name : exported_names) {

  const auto result_shape = result_type.getShape();
  // Index for the broadcasted matrix.
  llvm::SmallVector<int64_t, 16> current_index(result_type.getRank(), 0);
  // Computes the new operand shape using the original shape and the broadcast
  // dimensions to match result shape.
  llvm::SmallVector<int64_t, 16> operand_new_shape(result_type.getRank(), 1);
  for (int i = 0; i < dimensions.size(); ++i) {