return failure();
    }
  }

  return success();
}

LogicalResult CheckOutputConsumer(
    Operation* call_op, int expected_num_outputs,
    llvm::DenseSet<int> expected_consumer_indices) {
  const int num_results = call_op->getNumResults();
  if (num_results != expected_num_outputs) return failure();

  for (int i = 0; i < expected_num_outputs; ++i) {

  }
  auto a = Parameter(&builder, 0, ShapeUtil::MakeScalarShape(xla::F32), "a");
  auto b = Parameter(&builder, 1, ShapeUtil::MakeScalarShape(xla::F32), "b");
  XlaOp call_op = xla::Call(&builder, to_apply, {a, b});

  std::vector<XlaOp> tuple_values;
  tuple_values.push_back(call_op);
  xla::Tuple(&builder, tuple_values);

  TF_ASSERT_OK_AND_ASSIGN(XlaComputation computation, builder.Build());

          next_values_to_visit.push_back(use.getOwner()->getResult(0));
          continue;
        }

        if (auto call_op = llvm::dyn_cast<mlir::CallOpInterface>(owner)) {
          mlir::func::FuncOp func =
              llvm::dyn_cast<mlir::func::FuncOp>(call_op.resolveCallable());
          if (!func) continue;
          next_values_to_visit.push_back(
              func.getArgument(use.getOperandNumber()));

func.func private @tensor_invalid_1() -> !tfr.tensor<[N, T>

// -----

// expected-error@+1 {{unbalanced}}
func.func @tensor_invalid_2() -> !tfr.tensor<[N, T]

// -----

// CHECK-LABEL: call_op
func.func @call_op(%arg0: !tfr.tensor<T>, %arg1: !tfr.tensor_list<TL>, %arg2: i32) -> !tfr.tensor<K> {
  %0 = tfr.call @Foo(%arg0, %arg1, %arg2) : (!tfr.tensor<T>, !tfr.tensor_list<TL>, i32) -> !tfr.tensor<K>
  func.return %0 : !tfr.tensor<K>

// `ExtractSubgraphToFunc` adds exactly two "new" `Operations`, a FuncOp and
// a CallOp. Pass these back to the caller for setting more specific attributes
// after graph mutation has taken place.
struct OpsAdded {
  mlir::func::FuncOp func_op;
  mlir::func::CallOp call_op;
};

// Given a `Subgraph` containing a sequence of adjacent `Operations` from

      } else {
        return -1;
      }
    }
  };

  for (auto call_op :
       func.getRegion().getOps<TF::StatefulPartitionedCallOp>()) {
    for (int i = 0; i < call_op->getNumResults(); ++i) {
      if (IsResourceType(call_op->getResult(i).getType())) {
        Operation* symbol = symbol_table.lookup(call_op.getF());
        if (symbol == nullptr) {

// to set attributes specific to this pass.
void AddAttrs(OpsAdded& ops_added, OpBuilder& builder, int func_count) {
  func::FuncOp& added_func_op = ops_added.func_op;
  func::CallOp& added_call_op = ops_added.call_op;
  StringAttr interface_name =
      builder.getStringAttr(absl::StrCat("func_", func_count));

  added_func_op->setAttr(kInterfaceNameAttr, interface_name);

  op_builder.setInsertionPointToStart(&region.front());
  auto call_op = op_builder.create<mlir::func::CallOp>(
      loc, func.getFunctionType().getResults(), func.getSymName(),
      region.getArguments());
  op_builder.create<mlir::TFL::YieldOp>(loc, call_op.getResults());
}

void InlineStablehloOpRegion(mlir::Region& region, mlir::func::FuncOp func) {

      while (auto cast_op =
                 llvm::dyn_cast_or_null<CastOp>(value.getDefiningOp())) {
        // Consider cast compatibility in case
        //    %cast = "tf.Cast"(%0) : (tensor<2xi64>) -> tensor<2xf32>
        // is skipped.
        if (cast_op.getSrcT() != cast_op.getDstT()) {
          break;
        }
        value = cast_op.getOperand();
      }
      return value;
    };

  auto get_call_index = [&](mlir::Block& b) -> std::optional<int> {
    if (b.getOperations().size() != 2) return std::nullopt;
    if (auto call_op = dyn_cast<mlir::func::CallOp>(b.front()))
      return subgraph_index_map_.at(call_op.getCallee().str());
    return std::nullopt;
  };
  auto body_subgraph_index = get_call_index(op.getBody().front());