if (!buffer_ty.hasStaticShape()) {
      mlir::Value return_val = return_op->getOperand(idx);
      if (auto owner = mlir::dyn_cast_or_null<mlir::tensor::CastOp>(
              return_val.getDefiningOp())) {
        // For bounded dynamic type, get a static size by taking bounds as the
        // dimensions. These dimensions are marked as dynamic in xla::Shape
        // below.

  // If `value` is produced by tf.Dequantize op, then return the Dequantize op's
  // input. Otherwise return `value`.
  auto get_real_input_value = [](Value value) -> Value {
    Operation* defining_op = value.getDefiningOp();
    if (auto dequantize = dyn_cast_or_null<TF::DequantizeOp>(defining_op)) {
      return dequantize.getInput();
    } else if (auto dequantize =
                   dyn_cast_or_null<TFL::DequantizeOp>(defining_op)) {

  // Extend tuple if `value` is a tuple.
  // If `value` is an op result and the owner is a `mhlo.tuple`, simply unpack
  // the tuple.
  if (auto tuple_op = value.getDefiningOp<TupleOp>()) {
    auto tuple_operands = llvm::to_vector(tuple_op.getOperands());
    tuple_operands.push_back(token);
    return {create_tuple(tuple_operands)};
  }