if (IsTPUOp(input_value.get().getDefiningOp()) &&
          !backward_pass_ops.contains(input_value.get().getDefiningOp()))
        values_to_add_nodes.insert(input_value.get());
  }

  for (Value value : values_to_add_nodes) {
    builder.setInsertionPointAfter(value.getDefiningOp());
    std::vector<Type> types(num_replicas, value.getType());
    Location loc = value.getDefiningOp()->getLoc();

      resource_alias_analysis.GetResourceAliases(resource_iterator);

  auto is_generator = [](Value val) {
    if (mlir::isa<BlockArgument>(val)) return true;
    Operation* definition = val.getDefiningOp();
    return definition->getNumOperands() == 0 &&
           definition->getNumResults() == 1;
  };

  // Check all generator aliases (ops or function argument) are on CPU.

    // * the arg to the identity op is produced by a TPUExecuteOp.
    if (op_device && op_device.getValue().contains("TPU")) {
      return true;
    }

    Operation *def_op = arg.getDefiningOp();
    // If the arg to this identity op is the arg of a function, there's no
    // defining op.
    if (def_op != nullptr &&
        (isa<TF::TPUExecuteOp, TF::TPUExecuteAndUpdateVariablesOp>(def_op))) {

  output_types.reserve(num_return_operands);
  for (int i = 0; i != num_return_operands; ++i) {
    auto returned_value = terminator->getOperand(i);
    Operation* returned_op = returned_value.getDefiningOp();
    if (returned_op && returned_op->hasOneUse() &&
        llvm::isa<DequantizeOp>(returned_op)) {
      auto dequantize_op = llvm::cast<DequantizeOp>(returned_op);

      ResourceConstructingOps result(global_tensor);
      return result;
    }
  } else if (auto vh = dyn_cast<TF::VarHandleOp>(value.getDefiningOp())) {
    return ResourceConstructingOps(vh);
  } else if (auto it = dyn_cast<TF::IteratorOp>(value.getDefiningOp())) {
    return ResourceConstructingOps(it);
  }
  return ResourceConstructingOps();
}

ResourceConstructingOps ResourceConstructingOps::join(

  }

  if (v.getDefiningOp()->getNumOperands() == 1) {
    // If the value is originating from an unary op, assume it's something
    // simple like "cast" and keep tracing.
    return traceUpwardsToArgument(v.getDefiningOp()->getOperand(0), seen, out);
  } else {
    // Typically a tf.VarHandle op.
    return v.getDefiningOp()->emitOpError("Non constant predecessor");
  }
}

    Value resource_operand = assign_op.getOperand(0);
    Value assigned_value_operand = assign_op.getOperand(1);

    if (isa<TF::VarHandleOp>(resource_operand.getDefiningOp()) &&
        isa<TF::ConstOp>(assigned_value_operand.getDefiningOp())) {
      return success();
    } else {
      return failure();
    }
  }

  void rewrite(TF::AssignVariableOp assign_op,

  static std::optional<FoldAdaptor> Create(Operation* operation) {
    auto foldable_opr = [](Value val) -> bool {
      return !llvm::isa<BlockArgument>(val) &&
             llvm::isa<stablehlo::ConstantOp>(val.getDefiningOp());
    };
    if (!llvm::all_of(operation->getOperands(), foldable_opr)) {
      return std::nullopt;
    }
    return FoldAdaptor(operation);
  }

    mlir::TF::AssignVariableOp assign_var_op, BundleWriter& bundle_writer) {
  auto resource_operand = assign_var_op.getOperand(0);
  auto var_handle_op =
      llvm::dyn_cast<mlir::TF::VarHandleOp>(resource_operand.getDefiningOp());
  if (!var_handle_op) {
    assign_var_op->emitRemark(
        "Operand idx 0 is not a tf.VarHandleOp. The initializing tensor is not "
        "saved to checkpoint.");
    return "";
  }

                             Value multiplier) {
  auto dq_op = value.getDefiningOp<quantfork::DequantizeCastOp>();
  if (!dq_op) {
    auto mul_op = builder.create<TF::MulOp>(loc, value, multiplier);
    return mul_op.getResult();
  }
  auto q_op = dq_op.getArg().getDefiningOp<quantfork::QuantizeCastOp>();
  if (!q_op) return {};

  Value float_value = q_op.getArg();