// Propagate the analysis results from within the callee's body.
      PropagatePotentiallyWrittenUpFromCallee(batch_function.func().getRegion(),
                                              batch_function.getOperands());
      return;
    }
    // For all other ops, we assume it mutates all resources it uses, so
    // this errs on the side of being conservative. We should improve

    auto const_type = mlir::dyn_cast<ShapedType>(op.getOperand(0).getType());
    auto result_type = mlir::dyn_cast<ShapedType>(op.getResult().getType());
    if (!const_type || !const_type.getElementType().isSignlessInteger(4) ||
        !result_type || !result_type.getElementType().isSignlessInteger(8)) {
      return failure();
    }

    auto input_op =
        dyn_cast<arith::ConstantOp>(op.getOperand(0).getDefiningOp());
    if (!input_op) {

    Operation::operand_range mean(fused_batch_norm->getOperands());
    ::mlir::FloatAttr exponential_avg_factor;
    ::mlir::TF::FusedBatchNormV3Op root;
    Operation::operand_range offset(fused_batch_norm->getOperands());
    Operation::operand_range x(fused_batch_norm->getOperands());
    Operation::operand_range scale(fused_batch_norm->getOperands());
    Operation::operand_range variance(fused_batch_norm->getOperands());

  // This is the last op, should be tf_device::ReturnOp.
  assert(mlir::isa<mlir::tf_device::ReturnOp>(return_op));

  Value old_val, new_val;
  for (auto it : llvm::zip(cluster_op.getResults(), return_op.getOperands())) {
    std::tie(old_val, new_val) = it;
    old_val.replaceAllUsesExcept(new_val, &return_op);
  }

  cluster_op.erase();
}

void XlaInlineDeviceOpsPass::runOnOperation() {

      island.getBody().front().begin(), body.getOperations(),
      copy_range.begin(), copy_range.end());
  builder.setInsertionPointToEnd(&island.getBody().front());
  builder.create<tf_executor::YieldOp>(loc, return_op.getOperands());
  for (auto item : llvm::enumerate(graph_op.getResults())) {
    return_op.setOperand(item.index(), item.value());
  }
}

std::unique_ptr<OperationPass<func::FuncOp>>

}

}  // end anonymous namespace

OpFoldResult ReshapeOp::fold(FoldAdaptor adaptor) {
  auto operands = adaptor.getOperands();
  // Remove identity reshape with both static result and input shape.
  auto result_type = getType().cast<ShapedType>();
  auto input_type = getOperand(0).getType().cast<ShapedType>();
  if (InputOutputHasSameShape(input_type, result_type)) return getInput();

  // Constant folding

            mlir::cast<IntegerAttr>(mirrored_index).getInt());
      }
    }
    auto func =
        llvm::cast<func::FuncOp>(m.lookupSymbol(cluster_func.getFunc()));
    for (auto entry : llvm::enumerate(cluster_func.getOperands())) {
      auto operand = SkipIdentityAndReadVariable(entry.value());
      auto block_arg = mlir::dyn_cast<BlockArgument>(operand);
      if (block_arg && block_arg.getOwner() == &replicate.GetBody()) {

  auto old_body_ret = body.front().getTerminator();
  auto new_retvals = llvm::to_vector<8>(old_body_ret->getOperands());
  for (int64_t i = 0; i < while_op.getNumResults(); ++i) {
    if (!ta_arg_buffer_type(i)) continue;
    auto retval = old_body_ret->getOperand(i);
    auto arg = retval.dyn_cast<BlockArgument>();
    if (!arg) {
      return while_op.emitOpError(

      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 =
        FunctionType::get(ctx, func_operand_types, func_result_types);

    // Create the outlined function

  LogicalResult match(Operation* op) const override {
    if (isa<TF::CastOp, TF::ConstOp>(op) ||
        op->getName().hasTrait<OpTrait::ZeroOperands>()) {
      return failure();
    }
    for (Value input : op->getOperands()) {
      if (getElementTypeOrSelf(input).isBF16()) {
        return success();
      }
    }
    for (Value value : op->getResults()) {
      if (getElementTypeOrSelf(value).isBF16()) {
        return success();