const std::vector<std::string> indices =
            absl::StrSplit(node_specification, '-');
        for (const std::string& cur_index : indices) {
          custom_op_map[node_name].quantizable_input_indices.push_back(
              std::stoi(cur_index));
        }
        break;
      }
      case CustomOpUpdateOptions::kWeightOnly:
        custom_op_map[node_name].is_weight_only =
            GetBooleanSpecs(node_specification);

      sharding_for_args.push_back(operand_sharding);
      continue;
    }

    if (infer_from_computation) {
      auto arg_sharding = GetXlaShardingFromArg(arg, logical_device_vec);
      if (arg_sharding) {
        sharding_for_args.push_back(arg_sharding.value());
        continue;
      }
    }

    sharding_for_args.push_back(std::nullopt);
  }
}

  builder.setInsertionPoint(call);
  auto new_operands = llvm::to_vector(call.getArgOperands());
  new_operands.push_back(token);
  auto new_result_types = llvm::to_vector(call.getResultTypes());
  new_result_types.push_back(token.getType());
  auto new_call = builder.create<func::CallOp>(
      call.getLoc(), new_result_types,
      new_symbol ? *new_symbol : call.getCallee(), new_operands);

      // Add new control input if its defining op is not already a defining
      // op for some other operand. Update defining_ops.
      if (defining_ops.insert(new_control_input.getDefiningOp()).second) {
        operands.push_back(new_control_input);
      }
    }
    state.addOperands(operands);
    Operation* new_op = builder.create(state);
    item.replaceAllUsesWith(new_op);
    new_op->setAttrs(item.getAttrDictionary());

  if (!main_fn) return std::string("");

  flatbuffers::FlatBufferBuilder fb_builder;
  std::vector<mlir::func::FuncOp> funcs;
  funcs.push_back(main_fn);
  module.walk([&](mlir::func::FuncOp fn) {
    if (fn != main_fn) {
      funcs.push_back(fn);
    }
  });

  // Populate the hardware metadata.
  // And collect the hardwares used.
  std::map<std::string, uint8_t> hardware_map;

  output_tensor_names.reserve(fetch_outputs.size());
  for (auto const& output : fetch_outputs) {
    output_tensor_names.push_back(output.name());
  }
  std::vector<string> target_node_names;
  target_node_names.reserve(run_outputs.size());
  for (auto const& output : run_outputs) {
    target_node_names.push_back(output.node()->name());
  }
  TF_RETURN_IF_ERROR(impl()->MaybeExtendGraph());

  }
  auto iter = backprops_.find(src);
  if (iter != backprops_.end()) {
    auto* grads = &iter->second;
    grads->push_back(dst_grad);
    if (--pending_[src.node()->id()] == 0) {
      ready_.push_back(src.node());
    }
  }
  return absl::OkStatus();
}

std::vector<bool> SymbolicGradientBuilder::GetReachableNodes() {

    func_operand_types.reserve(operands.size());
    for (Value operand : operands)
      func_operand_types.push_back(operand.getType());

    // Function results are the yield operands
    SmallVector<Type, 16> func_result_types;
    for (Value operand : island_op.GetYield().getOperands())
      func_result_types.push_back(operand.getType());
    FunctionType func_type =

  // reshape back into a data-shaped tensor to propagate gradients for the data
  // input.
  grad_outputs->push_back(Reshape(scope, reconstructed, Shape(scope, data)));
  // Stop propagation along the partitions input
  grad_outputs->push_back(NoGradient());
  return scope.status();
}
REGISTER_GRADIENT_OP("DynamicPartition", DynamicPartitionGrad);

    } else {
      reduced_axes.push_back(*axis);
    }
  }
  // Get the corresponding dimensions for each reduced axis.
  std::vector<Output> reduced_dims_inputs;
  reduced_dims_inputs.reserve(reduced_axes.size());
  for (const int i : reduced_axes) {
    if (i < 0) {
      reduced_dims_inputs.push_back(
          Gather(scope, input_shape, Add(scope, Size(scope, input_shape), i)));