const std::vector<Output>& outputs,
                            const std::vector<Output>& inputs,
                            std::vector<Output>* grad_outputs) {
  std::vector<Output> grad_inputs;
  grad_inputs.reserve(outputs.size());
  for (const Output& output : outputs) {
    grad_inputs.emplace_back(ops::OnesLike(scope, output));
  }

  params_->body_outputs[1] = params_->body_inputs[1];

  // Finalize while loop
  ExpectOK();

  // Validate while loop outputs returned by TF_FinishWhile()
  EXPECT_TRUE(outputs_[0].oper != nullptr);
  EXPECT_GE(outputs_[0].index, 0);
  EXPECT_TRUE(outputs_[1].oper != nullptr);
  EXPECT_GE(outputs_[1].index, 0);

  // Check that cond and body inputs are not present
  for (int i = 0; i < params_->ninputs; ++i) {

          %outputs_2, %control_3 = tf_executor.island wraps "tf.XlaSpmdFullToShardShape"(%outputs_0) {dim = -1 : i64, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []} : (tensor<2x2xf32>) -> tensor<1x2xf32>

                @InputFiles def inputFiles = project.layout.files()

                @OutputDirectory File outputs1
                @OutputDirectory File outputs2

                @TaskAction void action() {
                    new File(outputs1, "output1.txt").text = "output1"
                    new File(outputs2, "output2.txt").text = "output2"
                }
            }
        """

        buildFile << """

John Howard <******@****.***> 1717524306 -0700

  const char* const code = R"mlir(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
  func.func @main() {
    %outputs_0 = "tf.TPUOrdinalSelector"() {device = ""} : () -> tensor<?xi32>
    "tf.TPUPartitionedCall"(%outputs_0) {f = @reachable_func} : (tensor<?xi32>) -> ()
    func.return
  }
  func.func @reachable_func() {
    func.return
  }
}
)mlir";