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");
  }
}

  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 "";
  }

  LogicalResult matchAndRewrite(quantfork::DequantizeCastOp op,
                                PatternRewriter& rewriter) const override {
    auto input_op = op.getArg().getDefiningOp();
    if (auto q = llvm::dyn_cast_or_null<quantfork::QuantizeCastOp>(input_op)) {
      if (!q->getAttr(kVolatileOpAttrName)) return failure();

      if (remove_volatile_ops_type == kPreserveInputsAndOutputs) {

def CloneF32ElementsAttrWithOnes
  : NativeCodeCall<"DenseElementsAttr::get($0.getType().cast<ShapedType>(), (float)1.0)">;

def NotConstant : Constraint<
    CPred<"$0.isa<BlockArgument>() || !llvm::isa<stablehlo::ConstantOp>($0.getDefiningOp())">,
    "Is not a constant.">;

def : Pat<(StableHLO_DivOp $l,
            (StableHLO_ConstantOp:$divisor FloatElementsAttr<32>:$cst)),
          (StableHLO_MulOp $l,
            (StableHLO_DivOp

  Operation* xla_call_module_op =
      FindOperationOfType<TF::XlaCallModuleOp>(main_fn);
  Operation* filter_dcast_op =
      xla_call_module_op->getOperand(1).getDefiningOp();
  Operation* filter_qcast_op = filter_dcast_op->getOperand(0).getDefiningOp();
  ASSERT_NE(filter_qcast_op, nullptr);
  EXPECT_TRUE(isa<quantfork::QuantizeCastOp>(filter_qcast_op));
  EXPECT_TRUE(isa<quantfork::DequantizeCastOp>(filter_dcast_op));

// returns the owner of the Block.
Operation* GetOpOfValue(Value value) {
  if (auto block_arg = mlir::dyn_cast<BlockArgument>(value))
    return block_arg.getOwner()->getParentOp();

  return value.getDefiningOp();
}

void TPUAnnotateDynamicShapeInputsPass::runOnOperation() {
  getOperation().walk([&](tf_device::ClusterFuncOp cluster_func_op) {
    Builder builder(cluster_func_op->getContext());

// Combines the two variadic arguments ($in_tensors and $captured_tensors).
def GetBatchFunctionOpArgOperands:
    NativeCodeCall<"cast<TF::BatchFunctionOp>($0[0].getDefiningOp()).getArgOperands()">;

// Replaces `TF_BatchFunctionOp` into `TF_PartitionedCallOp` that calls the
// same $f. This may be required, for example, when inlining is desired,

int64_t GetTransferredTensorBytes(func::CallOp from_graph,
                                  func::CallOp to_graph) {
  int64_t total_size_transferred = 0;
  for (auto input : to_graph.getOperands()) {
    Operation* input_op = input.getDefiningOp();
    if (input_op && input_op == from_graph.getOperation()) {
      auto input_type =
          mlir::dyn_cast_or_null<RankedTensorType>(input.getType());