for (auto i : return_indices[index_and_value.index()]) {
      merged_then_yield_values.push_back(
          if_op.getThenBranch().front().getTerminator()->getOperand(i));
    }
  }
  builder.setInsertionPointToEnd(&new_if_op.getThenBranch().front());
  builder.create<TF::YieldOp>(
      first_if.getThenBranch().front().getTerminator()->getLoc(),
      /*operands=*/merged_then_yield_values);

      builder.getStringAttr(if_op.else_function().getName()));
  CopyDeviceAndUnderscoredAttributes(if_op, if_region);

  CreateCall(if_op, if_op.then_function(),
             /*caller_region=*/if_region.getThenBranch(), if_op.getInput(),
             /*use_region_args=*/false,
             /*forward_block_args=*/false);
  CreateCall(if_op, if_op.else_function(),

  for (const auto& entry : output_buffer_to_size) {
    (*buffer_to_size)[new_op.getResult(std::get<0>(entry))] = {
        new_op.getResult(std::get<1>(entry)), std::get<2>(entry)};
  }

  new_op.getThenBranch().takeBody(if_op.getThenBranch());
  new_op.getElseBranch().takeBody(if_op.getElseBranch());

  if_op.replaceAllUsesWith(
      new_op.getResults().take_front(if_op.getNumResults()));
  if_op.erase();
  return success();
}

  }
  builder.create<TF::YieldOp>(loc,
                              /*operands=*/else_returns);

  // Then branch gets the data and calls the reduce_function.
  auto& then_branch = dataset_if.getThenBranch();
  then_branch.push_back(new Block);
  builder.setInsertionPointToEnd(&then_branch.front());
  // Add iterator operational data access inside if.

    // collect the call arguments post lookup through cast Op's
    extern_values.push_back(first);
    return true;
  };

  const TrivialTransformInfo tti(IsSingleCallRegion(if_region.getThenBranch()),
                                 IsSingleCallRegion(if_region.getElseBranch()),
                                 if_arg_matcher);

  std::string then_name, else_name;

  if (tti.can_transform) {

    }
  } else if (auto if_op = dyn_cast<TF::IfRegionOp>(user_op)) {
    (*arguments_to_erase)[if_op].push_back(argument_index);
    for (auto callee : {&if_op.getThenBranch(), &if_op.getElseBranch()}) {
      work_list->push_back(std::make_pair(callee, argument_index));
    }
  } else if (auto while_op = dyn_cast<TF::WhileOp>(user_op)) {