Method TestCopyConstantLayoutR4

tensorflow/compiler/xla/tests/copy_test.cc:188–220

}

void CopyOpTest::TestCopyConstantLayoutR4(size_t n1, size_t n2, size_t n3,
                                          size_t n4,
                                          absl::Span<const int64> permutation) {
  Array4D<int32> a(n1, n2, n3, n4);
  for (size_t i = 0; i < n1; ++i) {
    for (size_t j = 0; j < n2; ++j) {
      for (size_t k = 0; k < n3; ++k) {
        for (size_t l = 0; l < n4; ++l) {
          a(i, j, k, l) = i * n4 * n3 * n2 + j * n4 * n3 + k * n4 + l;
        }
      }
    }
  }

  HloComputation::Builder builder(TestName());

  Literal literal = LiteralUtil::CreateR4FromArray4D(a);

  HloInstruction* constant = builder.AddInstruction(
      HloInstruction::CreateConstant(std::move(literal)));

  builder.AddInstruction(HloInstruction::CreateUnary(
      constant->shape(), HloOpcode::kCopy, constant));

  std::unique_ptr<HloComputation> computation = builder.Build();

  auto module = CreateNewVerifiedModule();
  module->AddEntryComputation(std::move(computation));
  ForceResultLayout(module.get(), LayoutUtil::MakeLayout(permutation));
  Literal result = ExecuteAndTransfer(std::move(module), {});

  LiteralTestUtil::ExpectR4EqualArray4D(a, result);
}

XLA_TEST_F(CopyOpTest, CopyConstantR3Layout021_SingleIncompleteTilePerLayer) {
  TestCopyConstantLayout021(2, 2, 3);