}
  builder_.create<func::ReturnOp>(func_.getLoc(), ret_operands);

  auto arg_types = body.getArgumentTypes();
  auto result_types = body.getTerminator()->getOperandTypes();
  func_.setType(FunctionType::get(func_.getContext(), arg_types, result_types));
  *f = new MlirFunction(std::move(context_), std::move(module_), func_);
  return absl::OkStatus();
}

extern "C" {

// Updates func argument type to have the updated input shape.
void UpdateFuncType(func::FuncOp func) {
  auto arg_types = func.front().getArgumentTypes();
  auto result_types = func.front().getTerminator()->getOperandTypes();
  func.setType(FunctionType::get(func.getContext(), arg_types, result_types));
}

void HandleFuncOp(Operation* op) {
  auto func = llvm::cast<func::FuncOp>(op);
  UpdateFuncType(func);
}

    if (x_type != result_type) {
      x = builder.create<BroadcastInDimOp>(loc, result_type, x, broadcast_dims);
    }
    if (y_type != result_type) {
      y = builder.create<BroadcastInDimOp>(loc, result_type, y,
                                           larger_broadcast_dims);
    }
  } else {
    if (x_type != result_type) {
      x = builder.create<BroadcastInDimOp>(loc, result_type, x,

    SmallVector<Value> inputs{op.getInput()};
    SmallVector<Value> init_values{op.getInitValue()};
    SmallVector<Type> result_types{op.getResult().getType()};
    ReplaceTfOpWithNewOp<TF::XlaVariadicReduceV2Op>(
        rewriter, op, result_types, inputs, init_values,
        op.getDimensionsToReduce(), op.getReducer());
    return ::mlir::success();
  };
};

  // already set.
  auto symbol =
      mlir::SymbolRefAttr::get(builder.getContext(), func.getSymName());
  auto result_types = func.getResultTypes();
  auto caller = builder.create<TF::StatefulPartitionedCallOp>(
      loc, result_types, operands.getArrayRef(), symbol,
      /*config=*/builder.getStringAttr(""),
      /*config_proto=*/builder.getStringAttr(""),

    TypeRange operand_types, TypeRange result_types,
    TypeRange region_argument_types) {
  llvm::SmallVector<Type, 4> types;
  types.reserve(operand_types.size());
  for (auto entry :
       llvm::zip(operand_types, result_types, region_argument_types)) {
    auto operand_type = mlir::cast<TensorType>(std::get<0>(entry));
    auto result_type = mlir::cast<TensorType>(std::get<1>(entry));

mlir::TF::_XlaCompileMlirPlaceholderProgramKeyOp
CreateCompilationKeyPlaceholder(Location loc, OpBuilder& builder) {
  auto result_type =
      RankedTensorType::get({3}, builder.getType<mlir::TF::StringType>());
  return builder.create<mlir::TF::_XlaCompileMlirPlaceholderProgramKeyOp>(
      loc, /*program=*/result_type, llvm::ArrayRef<Value>{});
}

// Creates a `tf_device.launch` to wrap cluster ops.

  if (shape_value) {
    return rewriter.createOrFold<mhlo::DynamicBroadcastInDimOp>(
        loc, result_type, value1d, shape_value, dims);
  }
  assert(result_type.hasStaticShape());
  return rewriter.create<mhlo::BroadcastInDimOp>(loc, result_type, value1d,
                                                 dims);
}

// Gets the shape of operand, assuming it is a dynamic shape with static rank.

                           llvm::function_ref<APFloat(APFloat)> calculate) {
  assert(IsF32ShapedType(result_type) || IsBF16ShapedType(result_type));
  auto result_shape_type = result_type.cast<ShapedType>();

  if (!result_shape_type.hasStaticShape()) return {};

  if (auto dense_elements = operand.dyn_cast_or_null<DenseElementsAttr>()) {
    SmallVector<APFloat, 16> new_values;

  // result of 'resultType'. If a conversion can not be generated, nullptr
  // should be returned.
  Operation *materializeCallConversion(OpBuilder &builder, Value input,
                                       Type result_type,
                                       Location conversion_loc) const final {
    if (!mlir::isa<TensorType>(result_type) ||
        !mlir::isa<TensorType>(input.getType()))