getContext(), ftype.getResults(), graph_func.getAllResultAttrs(),
            res_types, res_attrs)))
      return failure();

    // Update the function type which has excluded the control args.
    func->setAttr("function_type", TypeAttr::get(rewriter.getFunctionType(
                                       arg_types, res_types)));

    // Update arg/result attributes.
    func.setAllArgAttrs(arg_attrs);

    new_config.push_back(builder.getNamedAttr(
        kCalledIndexAttrName, builder.getI64IntegerAttr(it->second)));

    // Set the `tf.backend_config` attribute to the `new_config`.
    op->setAttr(kTfBackendConfigAttrName,
                builder.getDictionaryAttr(new_config));

    return WalkResult::advance();
  });
  if (result.wasInterrupted()) {
    return failure();
  }

    return signalPassFailure();
  }
  ModuleOp outlined_module = ModuleOp::create(getOperation().getLoc());
  outlined_module->setAttrs(getOperation()->getAttrDictionary());
  outlined_module->setAttr(SymbolTable::getSymbolAttrName(),
                           StringAttr::get(ctx, kNestedModule));
  symbol_table.insert(outlined_module);
  SymbolTable outlined_symbol_table(outlined_module);

  // Jax wrapped the output nodes in a tuple, so it's pretty hard to us
  // to tell the output at this point, we will set the output at the export
  // phase.
  main_func->setAttr("tf.entry_function", builder.getDictionaryAttr(attrs));

  // StableHLO Quantizer is not supported for JAX input models, so
  // quantization_py_function_lib is set to nullptr.

              addtional_pad !=
          bts_crops_attr.getValues<APInt>()[{i, j}].getSExtValue()) {
        op->setAttr("padding", rewriter.getStringAttr("SAME"));
        break;
      }
    }
  }

  // Set dilations
  op->setAttr("dilations", dilations_attr.value());

  if (expand_op) {
    // If there is `expand_op`, we need to rewire the inputs to bypass the

      if (!matcher.Match(xla_call_module_op)) continue;

      // Set the text representation of `Method` to matched
      // `TF::XlaCallModuleOp`.
      xla_call_module_op->setAttr(
          kQuantizationMethodAttr,
          StringAttr::get(module_op.getContext(),
                          std::move(*quantization_method_txtpb)));
    }
  }
  return success();
}

template <typename Predicate>
void CopyAttributes(Operation *from, Operation *to, Predicate P) {
  for (const NamedAttribute &attr : from->getAttrs())
    if (P(attr)) to->setAttr(attr.getName(), attr.getValue());
}

// Copies attributes whose name begins with an _ from `from` to `to`.
inline void CopyUnderscoredAttributes(Operation *from, Operation *to) {

  }

  if (cifg) {
    attributes.push_back(kCoupleInputForgetGates);
  }

  mlir::StringAttr attr_values =
      builder->getStringAttr(llvm::join(attributes, ","));

  func->setAttr(kTFImplements, attr_values);
  return func;
}

class LstmUtilsTest : public ::testing::Test {
 protected:
  LstmUtilsTest() {}

  void SetUp() override {
    context_ = std::make_unique<mlir::MLIRContext>();

            values.push_back(range.second - range.first);
          }
          auto attr_value =
              mlir::DenseI32ArrayAttr::get(tf_op->getContext(), values);
          tf_op->setAttr(attr_name, attr_value);
        };
    if (tf_op->hasTrait<mlir::OpTrait::AttrSizedOperandSegments>() ||
        tf_op->hasTrait<mlir::OpTrait::AttrSizedResultSegments>()) {

  if (failed(::mlir::TF::UpdateDataFormat(data_format, this))) return failure();

  // Update convolution attributes.
  (*this)->setAttr("dilations", ShuffleArrayAttr(getDilations(), perm));
  (*this)->setAttr("strides", ShuffleArrayAttr(getStrides(), perm));
  (*this)->setAttr("explicit_paddings",
                   ShuffleArrayAttr(getExplicitPaddings(), perm, 2));

  return success();
}