for (auto& input : inputs) {
    auto input_op = input.get().getDefiningOp();
    if (!input_op || !IsSupportedHostInputOp(input_op)) return false;
  }
  for (auto entry : llvm::enumerate(inputs)) {
    Value input = entry.value().get();
    auto ranked_type = mlir::dyn_cast<RankedTensorType>(input.getType());
    if (!ranked_type) return false;
    auto input_shape = ranked_type.getShape();

    // Construct forward graph.
    auto a = Const(scope, 1, {4, 2});
    auto b = Const(scope, 2, {4, 2});
    auto c = Const(scope, 3, {4, 2});

    auto pack = Stack(scope, {a, b, c});
    auto unpack = Unstack(scope, pack.output, 3);
    TF_ASSERT_OK(scope.status());

    // Construct grad inputs.
    auto dx = Const(scope, 4, {4, 2});
    auto dy = Const(scope, 5, {4, 2});

  TensorShape scale_shape({shape.dim_size(channel_dim)});
  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
  auto scale = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(scale_shape));
  auto offset = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(scale_shape));
  auto mean = ops::ZerosLike(scope_, scale);
  auto var = ops::OnesLike(scope_, scale);

  if (!channel_first) {

    return rewriter.notifyMatchFailure(op, "zero_points must be constant");
  }

  auto original_element_type = getElementTypeOrSelf(original_type);
  if (!mlir::isa<TF::Qint8Type, TF::Qint32Type>(original_element_type)) {
    return rewriter.notifyMatchFailure(
        op, "Quantized type must be qint8 or qint32.");
  }
  auto storage_type = GetIntTypeFromTFQint(original_element_type);

  llvm::SmallPtrSet<Operation*, 4> defining_ops;
  llvm::SmallVector<Type, 4> types;
  for (auto& item :
       llvm::make_early_inc_range(llvm::reverse(graph_op.GetBody()))) {
    auto it = new_control_inputs.find(&item);
    if (it == new_control_inputs.end()) continue;
    auto& new_control_inputs_for_item = it->second;
    builder.setInsertionPoint(&item);

  // has the property that if there is an edge from SCC1->SCC2, SCC1 is visited
  // after SCC2, i.e., the graph is traversed bottom up just the way we want.
  auto scc_begin = llvm::scc_begin(&call_graph);
  auto scc_end = llvm::scc_end(&call_graph);
  for (auto& scc : make_range(scc_begin, scc_end)) {
    // Each SCC node is a collection of callgraph nodes that form a cycle. We

void RaiseTargetSubgraphsPass::runOnOperation() {
  ModuleOp module = getOperation();
  auto& side_effect_analysis = getAnalysis<TF::SideEffectAnalysis>();
  SmallVector<func::FuncOp> funcs(module.getOps<func::FuncOp>());
  int func_count = -1;
  for (auto func : funcs) {
    const auto& info = side_effect_analysis.GetAnalysisForFunc(func);
    for (auto& block : func) {
      OpBuilder builder = OpBuilder::atBlockBegin(&block);

  // maximal sharding then only the specified logical device take the value as
  // the input.
  for (const auto& sharding_attr_and_index : llvm::enumerate(sharding_attrs)) {
    const auto& sharding_attr = sharding_attr_and_index.value();
    const auto input_index = sharding_attr_and_index.index();
    const auto& input_value = cluster_func_inputs[input_index];

    xla::OpSharding sharding;
    if (DecodeShardingAttribute(

    void (*delete_func)(void*), TF_Status* status) {
  auto* cc_ctx = reinterpret_cast<tensorflow::OpKernelContext*>(ctx);
  auto* resource_mgr = cc_ctx->resource_manager();
  tensorflow::core::RefCountPtr<tensorflow::PluginResource>
      tf_plugin_resource_ptr;
  tensorflow::PluginResource* tf_plugin_resource = nullptr;

  auto cc_status = resource_mgr->LookupOrCreate<tensorflow::PluginResource>(

               std::vector<Output>* grad_outputs) {
  auto x = op.input(0);
  auto a = op.input(1);  // [Rank(x), 2]
  // Takes a slice of a. The 1st column. [Rank(x), 1].
  auto size = Stack(scope, {Rank(scope, x), 1});
  auto pad_before = Slice(scope, a, {0, 0}, size);
  // Make it a 1-D tensor.
  auto begin = Reshape(scope, pad_before, {-1});