Method TestCopyConstantLayout021

tensorflow/compiler/xla/tests/copy_test.cc:158–186

}

void CopyOpTest::TestCopyConstantLayout021(size_t n1, size_t n2, size_t n3) {
  Array3D<int32> a(n1, n2, n3);
  for (size_t i = 0; i < n1; ++i) {
    for (size_t j = 0; j < n2; ++j) {
      for (size_t k = 0; k < n3; ++k) {
        a(i, j, k) = i * n3 * n2 + j * n3 + k;
      }
    }
  }

  HloComputation::Builder builder(TestName());

  Literal literal = LiteralUtil::CreateR3FromArray3D(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({1, 2, 0}));
  Literal result = ExecuteAndTransfer(std::move(module), {});

  LiteralTestUtil::ExpectR3EqualArray3D(a, result);
}

void CopyOpTest::TestCopyConstantLayoutR4(size_t n1, size_t n2, size_t n3,
                                          size_t n4,