| 38 | class ClientTest : public ClientLibraryTestBase {}; |
| 39 | |
| 40 | XLA_TEST_F(ClientTest, ExecuteWithLayout) { |
| 41 | XlaBuilder b(TestName()); |
| 42 | |
| 43 | std::vector<std::vector<int64>> layouts = {{0, 1}, {1, 0}}; |
| 44 | for (const std::vector<int64>& execute_layout : layouts) { |
| 45 | for (const std::vector<int64>& transfer_layout : layouts) { |
| 46 | Add(ConstantR2<int32>(&b, {{1, 2}, {3, 4}}), |
| 47 | ConstantR2<int32>(&b, {{10, 20}, {30, 40}})); |
| 48 | TF_ASSERT_OK_AND_ASSIGN(auto computation, b.Build()); |
| 49 | |
| 50 | ExecutionOptions execution_options = execution_options_; |
| 51 | *execution_options.mutable_shape_with_output_layout() = |
| 52 | ShapeUtil::MakeShapeWithLayout(S32, /*dimensions=*/{2, 2}, |
| 53 | execute_layout) |
| 54 | .ToProto(); |
| 55 | TF_ASSERT_OK_AND_ASSIGN( |
| 56 | std::unique_ptr<GlobalData> data, |
| 57 | client_->Execute(computation, {}, &execution_options)); |
| 58 | |
| 59 | Literal expected_literal = LiteralUtil::CreateR2WithLayout<int32>( |
| 60 | {{11, 22}, {33, 44}}, LayoutUtil::MakeLayout(transfer_layout)); |
| 61 | |
| 62 | TF_ASSERT_OK_AND_ASSIGN( |
| 63 | auto computed, client_->Transfer(*data, &expected_literal.shape())); |
| 64 | |
| 65 | ASSERT_TRUE(LiteralTestUtil::EqualShapesAndLayouts( |
| 66 | expected_literal.shape(), computed.shape())); |
| 67 | EXPECT_TRUE(LiteralTestUtil::Equal(expected_literal, computed)); |
| 68 | } |
| 69 | } |
| 70 | } |
| 71 | |
| 72 | XLA_TEST_F(ClientTest, ExecuteWithTupleLayout) { |
| 73 | XlaBuilder b(TestName()); |
nothing calls this directly
no test coverage detected