Function XLA_TEST_F

tensorflow/compiler/xla/tests/while_test.cc:51–79

Tests a while node when the result type T is S32. int32 result = 0; while (result < 5) { result = result + 1; }

//   result = result + 1;
// }
XLA_TEST_F(WhileTest, WhileWithScalarS32Result) {
  auto result_shape = ShapeUtil::MakeShape(S32, {});

  // Create a computation for the condition: repeat for 5 iterations.
  XlaComputation condition;
  {
    XlaBuilder builder("condition");
    auto prev = Parameter(&builder, 0, result_shape, "prev");
    Gt(ConstantR0<int32>(&builder, 5), prev);
    condition = builder.Build().ConsumeValueOrDie();
  }

  // Create a computation for the body: add 1 to the result variable.
  XlaComputation body;
  {
    XlaBuilder builder("body");
    auto prev = Parameter(&builder, 0, result_shape, "prev");
    auto input = ConstantR0<int32>(&builder, 1);
    Add(input, prev);
    body = builder.Build().ConsumeValueOrDie();
  }

  // Create a While node with computations for the condition and the body.
  XlaBuilder builder(TestName());
  auto init = ConstantR0<int32>(&builder, 0);
  While(condition, body, init);

  ComputeAndCompareR0<int32>(&builder, 5, {});
}

// Tests a while node when the result type T is S64.
//