std::unordered_set<int> stop_backprop_nodes =
      GetStopBackpropNodes(reachable_nodes, output_nodes);

  // Populate `input_nodes_` from Outputs in `inputs_`.
  input_nodes_.reserve(inputs_.size());
  for (size_t i = 0; i < inputs_.size(); ++i) {
    input_nodes_.insert({inputs_[i], i});
  }

  // TODO(andydavis) Consider a more efficient data structure for `pending_` to

           const RunOptions& run_options = RunOptions(),
           RunMetadata* run_metadata = nullptr) {
    DCHECK_EQ(input_values.size(), inputs_.size());
    ClientSession::FeedType feeds;
    for (int i = 0; i < inputs_.size(); ++i) {
      feeds.emplace(inputs_[i], input_values[i]);
    }

    std::vector<Operation> run_outputs;
    std::vector<Tensor> out_tensors;

           const std::vector<T>& expected_output_values) {
    ClientSession session(scope_);

    DCHECK_EQ(input_values.size(), inputs_.size());
    ClientSession::FeedType feeds;
    for (int i = 0; i < inputs_.size(); ++i) {
      feeds.emplace(inputs_[i], input_values[i]);
    }

    std::vector<Tensor> out_tensors;
    TF_ASSERT_OK(session.Run(feeds, outputs_, &out_tensors));

  }

  void Run(const std::vector<TF_Output>& run_outputs,
           std::initializer_list<int> input_values) {
    DCHECK_EQ(inputs_.size(), input_values.size());
    std::vector<std::pair<TF_Operation*, TF_Tensor*>> inputs(inputs_.size());
    int i = 0;
    for (int v : input_values) {
      inputs[i] = {inputs_[i].oper, Int32Tensor(v)};
      ++i;
    }
    // TODO(skyewm): use std::make_unique or absl::make_unique when possible.

  mlir::MLIRContext* context = tfr_module_.getContext();
  llvm::SmallVector<mlir::Type, 4> input_tys, output_tys;
  mlir::Builder builder(context);
  for (auto ty : input_dtys) {
    mlir::Type elt_ty;
    TF_RETURN_IF_ERROR(ConvertDataType(ty, builder, &elt_ty));
    mlir::TensorType mlir_ty = mlir::UnrankedTensorType::get(elt_ty);
    input_tys.push_back(mlir_ty);
  }
  for (auto ty : output_dtys) {
    mlir::Type elt_ty;

    // i.e., num_input/input_last_dim = num_result/k
    auto input_ty = mlir::dyn_cast_or_null<ShapedType>(value.getType());
    auto result_ty = mlir::dyn_cast<ShapedType>(slice_op.getType());
    if (!input_ty || !result_ty) return std::nullopt;
    if (!input_ty.hasStaticShape() || !result_ty.hasStaticShape())
      return std::nullopt;
    if (!input_ty.getRank() || !result_ty.getRank()) return std::nullopt;

      }
      graph_outputs.push_back(output->output_);
    }

    auto s = TF_NewStatus();
    auto func = TF_GraphToFunction(graph_.get(), name_.data(), 0, -1, nullptr,
                                   inputs_.size(), inputs_.data(),
                                   graph_outputs.size(), graph_outputs.data(),
                                   nullptr, nullptr, name_.data(), s);
    *f = new GraphFunction(std::move(func->record->fdef()));

  const RankedTensorType input_ty =
      input.getType().dyn_cast_or_null<RankedTensorType>();
  const RankedTensorType filter_ty =
      filter.getType().dyn_cast_or_null<RankedTensorType>();
  // If indeed both input type & filter type are ranked type and have ranks.
  // We will need to check their ranks are valid.
  if ((input_ty && input_ty.hasRank() && input_ty.getRank() != 4) ||

  TF_EXPECT_OK(InitOp());

  // Add inputs.
  Tensor* a = CreateDeviceTensor<int32>(TensorShape({1, 3}), {1, 2, 3});
  Tensor* b = CreateDeviceTensor<int32>(TensorShape({1, 3}), {4, 5, 6});
  inputs_.push_back({nullptr, a});
  inputs_.push_back({nullptr, b});

  CreateContext();

  std::vector<XlaCompiler::Argument> args(2);
  args[0].kind = XlaCompiler::Argument::kParameter;
  args[0].type = DT_INT32;

    private final ImmutableMap<E, Integer> inputSet;
    private final int mask;

    SubSet(ImmutableMap<E, Integer> inputSet, int mask) {
      this.inputSet = inputSet;
      this.mask = mask;
    }

    @Override
    public Iterator<E> iterator() {
      return new UnmodifiableIterator<E>() {
        final ImmutableList<E> elements = inputSet.keySet().asList();
        int remainingSetBits = mask;