liveouts.snapshot_previous_state();
      return;
    }

    reverse_subgraph.insert(op);
    defined_values.insert(op->getResults().begin(), op->getResults().end());
    operands.insert(op->getOperands().begin(), op->getOperands().end());
  };

  for (Operation* op : reverse_main_func_block_ops) {
    if (!ops_to_add.contains(op)) continue;
    // When hitting a non-StableHLO op, i.e. tf.CustomAggregatorOp, start

    return std::nullopt;
  for (auto [arg, operand] :
       llvm::zip(block.getArguments(), call.getOperands())) {
    if (arg != operand) return std::nullopt;
  }
  for (auto [ret, operand] :
       llvm::zip(call.getResults(), yield.getOperands())) {
    if (ret != operand) return std::nullopt;
  }
  SymbolRefAttr symbol = call.getCallableForCallee().get<SymbolRefAttr>();

      // 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());

}

}  // 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()) {

      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

  if (!graph) return false;

  Operation* terminator = block.getTerminator();
  if (graph.getNumResults() != terminator->getNumOperands()) return false;
  for (auto result : llvm::zip(graph.getResults(), terminator->getOperands()))
    if (std::get<0>(result) != std::get<1>(result)) return false;

  return true;
}

// Checks if an operation of the tf_executor dialect can have TPU devices
// propagated through.

                 OpBuilder &builder) {
  llvm::SmallVector<Value> non_resource_args, resource_args;
  bool has_resources = false, in_order = true;
  for (const Value &arg : cluster_func_op.getOperands()) {
    if (!mlir::isa<TF::ResourceType>(getElementTypeOrSelf(arg.getType()))) {
      non_resource_args.push_back(arg);
      if (has_resources) in_order = false;
    } else {
      resource_args.push_back(arg);