experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Propagate test error to avoid crashes.

Browse Source 
PiperOrigin-RevId: 331853257
Change-Id: Idea8c0a4855b0c404f6ff6bd042ce950e2904767
This commit is contained in:
A. Unique TensorFlower 2020-09-15 14:17:48 -07:00 committed by TensorFlower Gardener
parent e8384561f1
commit 5cbcf59ba2
6 changed files with 182 additions and 124 deletions

15
tensorflow/compiler/xla/client/lib/slicing_test.cc
View File

@ -206,10 +206,12 @@ XLA_TEST_F(SlicingTest, DoubleEmptyIndexSelect) {
xla::XlaOp input, index;
Literal l(ShapeUtil::MakeShape(F32, {0, 1, 2, 0}));
Literal i(ShapeUtil::MakeShape(S32, {0}));
auto input_data =
CreateParameterAndTransferLiteral(0, l, "input", &builder, &input);
auto index_data =
CreateParameterAndTransferLiteral(1, i, "index", &builder, &index);
TF_ASSERT_OK_AND_ASSIGN(
auto input_data,
CreateParameterAndTransferLiteral(0, l, "input", &builder, &input));
TF_ASSERT_OK_AND_ASSIGN(
auto index_data,
CreateParameterAndTransferLiteral(1, i, "index", &builder, &index));
TorchIndexSelect(input, index, 0);
ComputeAndCompareLiteral(&builder, l, {input_data.get(), index_data.get()});
}
@ -219,8 +221,9 @@ XLA_TEST_F(SlicingTest, EmptyIndexSelectNonZero) {
xla::XlaOp input, index;
Literal l(ShapeUtil::MakeShape(F32, {0, 2}));
auto input_data =
CreateParameterAndTransferLiteral(0, l, "input", &builder, &input);
TF_ASSERT_OK_AND_ASSIGN(
auto input_data,
CreateParameterAndTransferLiteral(0, l, "input", &builder, &input));
auto index_data =
CreateR1Parameter<int>({0, 0, 0}, 1, "index", &builder, &index);
TorchIndexSelect(input, index, 0);

9
tensorflow/compiler/xla/tests/client_library_test_base.cc
View File

@ -605,7 +605,7 @@ XlaOp ClientLibraryTestBase::CreateConstantFromLiteral(const Literal& literal,
: LiteralSlice(literal));
}
std::unique_ptr<GlobalData>
StatusOr<std::unique_ptr<GlobalData>>
ClientLibraryTestBase::CreateParameterAndTransferLiteral(int64 parameter_number,
const Literal& literal,
const string& name,
@ -637,15 +637,14 @@ Literal ClientLibraryTestBase::MaybeConvertLiteralToBfloat16(
return literal.Clone();
}
std::unique_ptr<GlobalData>
StatusOr<std::unique_ptr<GlobalData>>
ClientLibraryTestBase::CreateParameterAndTransferLiteral(
int64 parameter_number, const Literal& literal, const string& name,
const DeviceHandle* device_handle, XlaBuilder* builder,
XlaOp* data_handle) {
Literal param_literal = MaybeConvertLiteralToBfloat16(literal);
std::unique_ptr<GlobalData> data =
client_->TransferToServer(param_literal, device_handle)
.ConsumeValueOrDie();
TF_ASSIGN_OR_RETURN(auto data,
client_->TransferToServer(param_literal, device_handle));
*data_handle =
Parameter(builder, parameter_number, param_literal.shape(), name);
return data;

4
tensorflow/compiler/xla/tests/client_library_test_base.h
View File

@ -270,14 +270,14 @@ class ClientLibraryTestBase : public ManifestCheckingTest {
// server, then stores into "data_handle" the global handle for that
// parameter. When the use_bfloat16 flag is set but the literal has F32
// elements, the literal will be converted to BF16 before being transferred.
std::unique_ptr<GlobalData> CreateParameterAndTransferLiteral(
StatusOr<std::unique_ptr<GlobalData>> CreateParameterAndTransferLiteral(
int64 parameter_number, const Literal& literal, const string& name,
XlaBuilder* builder, XlaOp* data_handle);
// As above, but the caller can specify the device that the literal is
// transferred to. If device_handle is nullptr, the literal will be
// transferred to the default device.
std::unique_ptr<GlobalData> CreateParameterAndTransferLiteral(
StatusOr<std::unique_ptr<GlobalData>> CreateParameterAndTransferLiteral(
int64 parameter_number, const Literal& literal, const string& name,
const DeviceHandle* device_handle, XlaBuilder* builder,
XlaOp* data_handle);

14
tensorflow/compiler/xla/tests/dot_operation_test.cc
View File

@ -69,12 +69,14 @@ XLA_TEST_F(DotOperationTest, DotOfInputTupleElem) {
XlaBuilder builder(TestName());
XlaOp param;
auto param_data = CreateParameterAndTransferLiteral(
0,
LiteralUtil::MakeTupleFromSlices(
{LiteralUtil::CreateR2<float>({{1, 2}, {3, 4}}),
LiteralUtil::CreateR2<float>({{5, 6}, {7, 8}})}),
"arg0", &builder, &param);
TF_ASSERT_OK_AND_ASSIGN(
auto param_data,
CreateParameterAndTransferLiteral(
0,
LiteralUtil::MakeTupleFromSlices(
{LiteralUtil::CreateR2<float>({{1, 2}, {3, 4}}),
LiteralUtil::CreateR2<float>({{5, 6}, {7, 8}})}),
"arg0", &builder, &param));
auto lhs = GetTupleElement(param, 0);
auto rhs = GetTupleElement(param, 1);
Dot(lhs, rhs);

34
tensorflow/compiler/xla/tests/reduce_window_test.cc
View File

@ -365,8 +365,9 @@ XLA_TEST_P(ReduceWindowTest, R4UnitWindow) {
Literal input_literal = LiteralUtil::CreateR4FromArray4DWithLayout(
input_array, LayoutUtil::MakeLayout({0, 3, 2, 1}));
XlaOp input;
auto input_data = CreateParameterAndTransferLiteral(
0, input_literal, "parameter", &builder_, &input);
TF_ASSERT_OK_AND_ASSIGN(
auto input_data, CreateParameterAndTransferLiteral(
0, input_literal, "parameter", &builder_, &input));
Padding padding = Padding::kSame;
ReduceWindowAdd(input, {1, 1, 7, 1}, {1, 4, 1, 1}, padding);
@ -423,8 +424,9 @@ XLA_TEST_P(ReduceWindowTest, R4SecondMinorStride) {
Literal input_literal = LiteralUtil::CreateR4FromArray4DWithLayout(
input_array, LayoutUtil::MakeLayout({3, 2, 1, 0}));
XlaOp input;
auto input_data = CreateParameterAndTransferLiteral(
0, input_literal, "parameter", &builder_, &input);
TF_ASSERT_OK_AND_ASSIGN(
auto input_data, CreateParameterAndTransferLiteral(
0, input_literal, "parameter", &builder_, &input));
int win_len = 1;
int stride = 8;
@ -444,8 +446,9 @@ XLA_TEST_P(ReduceWindowTest, R4SecondMinorUnitStride) {
Literal input_literal = LiteralUtil::CreateR4FromArray4DWithLayout(
input_array, LayoutUtil::MakeLayout({3, 2, 1, 0}));
XlaOp input;
auto input_data = CreateParameterAndTransferLiteral(
0, input_literal, "parameter", &builder_, &input);
TF_ASSERT_OK_AND_ASSIGN(
auto input_data, CreateParameterAndTransferLiteral(
0, input_literal, "parameter", &builder_, &input));
int win_len = 3;
int stride = 1;
@ -465,8 +468,9 @@ XLA_TEST_P(ReduceWindowTest, R4SecondMinorWin) {
Literal input_literal = LiteralUtil::CreateR4FromArray4DWithLayout(
input_array, LayoutUtil::MakeLayout({3, 2, 1, 0}));
XlaOp input;
auto input_data = CreateParameterAndTransferLiteral(
0, input_literal, "parameter", &builder_, &input);
TF_ASSERT_OK_AND_ASSIGN(
auto input_data, CreateParameterAndTransferLiteral(
0, input_literal, "parameter", &builder_, &input));
int win_len = 8;
int stride = 5;
@ -631,8 +635,9 @@ class R4ReduceWindowTest : public ReduceWindowTestBase,
Literal input_literal = LiteralUtil::CreateR4FromArray4DWithLayout(
input, LayoutUtil::MakeLayout(param.layout));
XlaOp parameter;
auto input_arg = CreateParameterAndTransferLiteral(0, input_literal, "p0",
&b, &parameter);
TF_ASSERT_OK_AND_ASSIGN(auto input_arg,
CreateParameterAndTransferLiteral(
0, input_literal, "p0", &b, &parameter));
std::vector<std::pair<int64, int64>> padding(4);
for (int i = 0; i < 4; ++i) {
@ -1243,7 +1248,9 @@ class R2ReduceWindowTest : public ReduceWindowTestBase,
input, LayoutUtil::MakeLayout(param.layout));
XlaOp parameter;
CreateParameterAndTransferLiteral(0, input_literal, "p0", &b, &parameter);
TF_ASSERT_OK(CreateParameterAndTransferLiteral(0, input_literal, "p0", &b,
&parameter)
.status());
std::vector<std::pair<int64, int64>> padding(2);
for (int i = 0; i < 2; ++i) {
@ -1443,8 +1450,9 @@ XLA_TEST_P(R1ReduceWindowTest, DoIt) {
Literal input_literal =
LiteralUtil::CreateR1(absl::Span<const float>(input_vector));
XlaOp parameter;
auto input_arg =
CreateParameterAndTransferLiteral(0, input_literal, "p0", &b, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input_arg, CreateParameterAndTransferLiteral(0, input_literal, "p0",
&b, &parameter));
std::vector<std::pair<int64, int64>> padding(1);
padding[0] = {param.pad_low[0], param.pad_high[0]};

230
tensorflow/compiler/xla/tests/reshape_test.cc
View File

@ -57,8 +57,9 @@ XLA_TEST_P(ReshapeTest, CollapseTrivial1x1) {
input_array.Fill(1.0f);
auto input_literal = LiteralUtil::CreateR2FromArray2D(input_array);
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "parameter",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(
0, input_literal, "parameter", &builder, &parameter));
Collapse(/*operand=*/parameter, /*dimensions=*/{0, 1});
auto expected_literal = LiteralUtil::CreateR1<float>({1.0f});
@ -70,8 +71,9 @@ XLA_TEST_P(ReshapeTest, CollapseTrivialR1EmptyDims) {
XlaBuilder builder(TestName());
auto input_literal = LiteralUtil::CreateR1<float>({1.0f});
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "parameter",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(
0, input_literal, "parameter", &builder, &parameter));
Collapse(/*operand=*/parameter, /*dimensions=*/{});
auto expected_literal = LiteralUtil::CreateR1<float>({1.0f});
@ -83,8 +85,9 @@ XLA_TEST_P(ReshapeTest, CollapseTrivialR1OnlyDim) {
XlaBuilder builder(TestName());
auto input_literal = LiteralUtil::CreateR1<float>({1.0f});
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "parameter",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(
0, input_literal, "parameter", &builder, &parameter));
Collapse(/*operand=*/parameter, /*dimensions=*/{0});
auto expected_literal = LiteralUtil::CreateR1<float>({1.0f});
@ -99,8 +102,9 @@ XLA_TEST_P(ReshapeTest, SingleElementArrayToScalar) {
input_array.Fill(1.0f);
auto input_literal = LiteralUtil::CreateR2FromArray2D(input_array);
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "parameter",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(
0, input_literal, "parameter", &builder, &parameter));
auto reshape = Reshape(/*operand=*/parameter, /*dimensions=*/{0, 1},
/*new_sizes=*/{});
auto new_shape = builder.GetShape(reshape).ConsumeValueOrDie();
@ -115,8 +119,9 @@ XLA_TEST_P(ReshapeTest, ScalarToSingleElementArray) {
Literal param0_literal = LiteralUtil::CreateR0<float>(1.0f);
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, param0_literal, "param0",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, param0_literal, "param0",
&builder, &parameter));
auto a = Neg(parameter);
Reshape(/*operand=*/a, /*dimensions=*/{}, /*new_sizes=*/{1});
@ -130,8 +135,9 @@ XLA_TEST_P(ReshapeTest, Trivial0x3) {
Array2D<float> input_array(0, 3);
auto input_literal = LiteralUtil::CreateR2FromArray2D(input_array);
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Collapse(/*operand=*/parameter, /*dimensions=*/{0, 1});
auto expected_literal = LiteralUtil::CreateR1<float>({});
ComputeAndCompareLiteral(&builder, expected_literal, {input.get()},
@ -144,8 +150,9 @@ XLA_TEST_P(ReshapeTest, Trivial0x3WithParameter) {
Literal param0_literal =
LiteralUtil::CreateR2FromArray2D<float>(Array2D<float>(0, 3));
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, param0_literal, "param0",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, param0_literal, "param0",
&builder, &parameter));
Collapse(/*operand=*/parameter, /*dimensions=*/{0, 1});
auto expected_literal = LiteralUtil::CreateR1<float>({});
ComputeAndCompareLiteral(&builder, expected_literal, {input.get()},
@ -157,8 +164,9 @@ XLA_TEST_P(ReshapeTest, Trivial3x0) {
Array2D<float> input_array(3, 0);
auto input_literal = LiteralUtil::CreateR2FromArray2D(input_array);
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Collapse(/*operand=*/parameter, /*dimensions=*/{0, 1});
auto expected_literal = LiteralUtil::CreateR1<float>({});
ComputeAndCompareLiteral(&builder, expected_literal, {input.get()},
@ -170,8 +178,9 @@ XLA_TEST_P(ReshapeTest, Trivial1x3) {
XlaBuilder builder(TestName());
auto input_literal = LiteralUtil::CreateR2<float>({{1.0f, 2.0f, 3.0f}});
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Collapse(/*operand=*/parameter, /*dimensions=*/{0, 1});
auto expected_literal = LiteralUtil::CreateR1<float>({1.0f, 2.0f, 3.0f});
ComputeAndCompareLiteral(&builder, expected_literal, {input.get()},
@ -183,8 +192,9 @@ XLA_TEST_P(ReshapeTest, Trivial3x1) {
XlaBuilder builder(TestName());
auto input_literal = LiteralUtil::CreateR2<float>({{1.0f}, {2.0f}, {3.0f}});
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Collapse(/*operand=*/parameter, /*dimensions=*/{0, 1});
auto expected_literal = LiteralUtil::CreateR1<float>({1.0f, 2.0f, 3.0f});
ComputeAndCompareLiteral(&builder, expected_literal, {input.get()},
@ -196,8 +206,9 @@ XLA_TEST_P(ReshapeTest, R1ToR2_0_To_2x0) {
XlaBuilder builder(TestName());
auto input_literal = LiteralUtil::CreateR1<float>({});
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Reshape(/*operand=*/parameter, /*dimensions=*/{0},
/*new_sizes=*/{2, 0});
auto expected_literal = LiteralUtil::CreateR2<float>({{}, {}});
@ -211,8 +222,9 @@ XLA_TEST_P(ReshapeTest, R1ToR2_6_To_2x3) {
auto input_literal =
LiteralUtil::CreateR1<float>({1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f});
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Reshape(/*operand=*/parameter, /*dimensions=*/{0},
/*new_sizes=*/{2, 3});
auto expected_literal =
@ -226,8 +238,9 @@ XLA_TEST_P(ReshapeTest, Reshape0x2To2x0) {
XlaBuilder builder(TestName());
auto input_literal = LiteralUtil::CreateFromArray(Array2D<float>(0, 2));
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Reshape(/*operand=*/parameter, /*dimensions=*/{0, 1},
/*new_sizes=*/{2, 0});
auto expected_literal = LiteralUtil::CreateR2<float>({{}, {}});
@ -241,8 +254,9 @@ XLA_TEST_P(ReshapeTest, ReshapeRowToCol) {
auto simple = MakeLinspaceArray2D(1.0f, 3.0f, 1, 3);
auto input_literal = LiteralUtil::CreateFromArray(*simple);
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Reshape(/*operand=*/parameter, /*dimensions=*/{0, 1},
/*new_sizes=*/{3, 1});
@ -258,8 +272,9 @@ XLA_TEST_P(ReshapeTest, TransposeAsReshape) {
auto a4x3 = MakeLinspaceArray2D(1.0f, 12.0f, 4, 3);
auto input_literal = LiteralUtil::CreateFromArray(*a4x3);
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Reshape(/*operand=*/parameter, /*dimensions=*/{1, 0},
/*new_sizes=*/{3, 4});
@ -274,8 +289,9 @@ XLA_TEST_P(ReshapeTest, Transpose0x4) {
XlaBuilder builder(TestName());
auto input_literal = LiteralUtil::CreateFromArray(Array2D<float>(0, 4));
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Transpose(parameter, {1, 0});
auto expected_literal = LiteralUtil::CreateR2<float>({{}, {}, {}, {}});
ComputeAndCompareLiteral(&builder, expected_literal, {input.get()},
@ -288,8 +304,9 @@ XLA_TEST_P(ReshapeTest, Transpose4x3) {
auto a4x3 = MakeLinspaceArray2D(1.0f, 12.0f, 4, 3);
auto input_literal = LiteralUtil::CreateFromArray(*a4x3);
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Transpose(parameter, {1, 0});
auto expected = ReferenceUtil::TransposeArray2D(*a4x3);
@ -304,8 +321,9 @@ XLA_TEST_P(ReshapeTest, ReshapeSplitNoShuffleZeroElements) {
XlaBuilder builder(TestName());
auto input_literal = LiteralUtil::CreateFromArray(Array2D<float>(6, 0));
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Reshape(/*operand=*/parameter, /*dimensions=*/{0, 1},
/*new_sizes=*/{2, 3, 0, 0});
auto expected_literal =
@ -318,8 +336,9 @@ XLA_TEST_P(ReshapeTest, ReshapeR4ToR2ZeroElements) {
XlaBuilder builder(TestName());
auto input_literal = LiteralUtil::CreateFromArray(Array4D<float>(2, 3, 4, 0));
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Reshape(/*operand=*/parameter, /*dimensions=*/{0, 1, 2, 3},
/*new_sizes=*/{24, 0});
auto expected_literal = LiteralUtil::CreateFromArray(Array2D<float>(24, 0));
@ -334,8 +353,9 @@ XLA_TEST_P(ReshapeTest, ReshapeSplitNoShuffle) {
auto a4x3 = MakeLinspaceArray2D(1.0f, 12.0f, 4, 3);
auto input_literal = LiteralUtil::CreateFromArray(*a4x3);
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Reshape(/*operand=*/parameter, /*dimensions=*/{0, 1},
/*new_sizes=*/{2, 6});
@ -349,8 +369,9 @@ XLA_TEST_P(ReshapeTest, ReshapeSplitAndShuffleZeroElements) {
XlaBuilder builder(TestName());
auto input_literal = LiteralUtil::CreateFromArray(Array2D<float>(0, 6));
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Reshape(/*operand=*/parameter, /*dimensions=*/{1, 0},
/*new_sizes=*/{3, 0});
auto expected_literal = LiteralUtil::CreateFromArray(Array2D<float>(3, 0));
@ -365,8 +386,9 @@ XLA_TEST_P(ReshapeTest, ReshapeSplitAndShuffle) {
auto a4x3 = MakeLinspaceArray2D(1.0f, 12.0f, 4, 3);
auto input_literal = LiteralUtil::CreateFromArray(*a4x3);
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Reshape(/*operand=*/parameter, /*dimensions=*/{1, 0},
/*new_sizes=*/{2, 6});
Array2D<float> expected({{1.0f, 4.0f, 7.0f, 10.0f, 2.0f, 5.0f},
@ -391,8 +413,9 @@ XLA_TEST_P(ReshapeTest, DocR3_R1_Collapse_012) {
XlaBuilder builder(TestName());
auto input_literal = LiteralUtil::CreateFromArray(ArrayForDocR3Tests());
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Reshape(/*operand=*/parameter, /*dimensions=*/{0, 1, 2},
/*new_sizes=*/{24});
auto expected_literal = LiteralUtil::CreateR1<float>(
@ -406,8 +429,9 @@ XLA_TEST_P(ReshapeTest, DocR3_R2_Collapse_012_Refine_83) {
XlaBuilder builder(TestName());
auto input_literal = LiteralUtil::CreateFromArray(ArrayForDocR3Tests());
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Reshape(/*operand=*/parameter, /*dimensions=*/{0, 1, 2},
/*new_sizes=*/{8, 3});
auto expected_literal = LiteralUtil::CreateR2<float>({{10, 11, 12},
@ -426,8 +450,9 @@ XLA_TEST_P(ReshapeTest, DocR3_R1_Collapse_120) {
XlaBuilder builder(TestName());
auto input_literal = LiteralUtil::CreateFromArray(ArrayForDocR3Tests());
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Reshape(/*operand=*/parameter, /*dimensions=*/{1, 2, 0},
/*new_sizes=*/{24});
auto expected_literal = LiteralUtil::CreateR1<float>(
@ -441,8 +466,9 @@ XLA_TEST_P(ReshapeTest, DocR3_R2_Collapse_120_Refine_83) {
XlaBuilder builder(TestName());
auto input_literal = LiteralUtil::CreateFromArray(ArrayForDocR3Tests());
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Reshape(/*operand=*/parameter, /*dimensions=*/{1, 2, 0},
/*new_sizes=*/{8, 3});
auto expected_literal = LiteralUtil::CreateR2<float>({{10, 20, 30},
@ -461,8 +487,9 @@ XLA_TEST_P(ReshapeTest, DocR3_R3_Collapse_120_Refine_262) {
XlaBuilder builder(TestName());
auto input_literal = LiteralUtil::CreateFromArray(ArrayForDocR3Tests());
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Reshape(/*operand=*/parameter, /*dimensions=*/{1, 2, 0},
/*new_sizes=*/{2, 6, 2});
auto expected_literal = LiteralUtil::CreateR3<float>(
@ -494,8 +521,9 @@ XLA_TEST_P(ReshapeTest, FullyConnectedCollapse) {
t2x2x2x3.FillWithYX(*filler2x3);
auto input_literal = LiteralUtil::CreateFromArray(t2x2x2x3);
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Collapse(/*operand=*/parameter, /*dimensions=*/{1, 2, 3});
auto expected_literal = LiteralUtil::CreateR2<float>(
{{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f},
@ -519,8 +547,9 @@ XLA_TEST_P(ReshapeTest, FullyConnectedCollapseDesugared) {
t(1, 0, 1, 1) = 7;
auto input_literal = LiteralUtil::CreateFromArray(t);
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Reshape(/*operand=*/parameter, /*dimensions=*/{0, 1, 2, 3},
/*new_sizes=*/{2, 4});
@ -542,8 +571,9 @@ XLA_TEST_P(ReshapeTest, ToScalar) {
input_literal.Set<float>(zeros, 83.0f);
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&b, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&b, &parameter));
Reshape(parameter, dimensions, {});
auto expected_literal = LiteralUtil::CreateR0<float>(83.0f);
@ -556,8 +586,9 @@ XLA_TEST_P(ReshapeTest, BadDimensions) {
XlaBuilder b(TestName());
auto input_literal = LiteralUtil::CreateR1<float>({1.0f});
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input", &b,
&parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&b, &parameter));
Reshape(parameter, {}, {});
EXPECT_THAT(
ExecuteToString(&b, {}),
@ -568,8 +599,9 @@ XLA_TEST_P(ReshapeTest, BadNewSizes) {
XlaBuilder b(TestName());
auto input_literal = LiteralUtil::CreateR1<float>({1.0f, 2.0f});
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input", &b,
&parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&b, &parameter));
Reshape(parameter, {1}, {});
EXPECT_THAT(ExecuteToString(&b, {}),
::testing::HasSubstr("mismatched element counts"));
@ -604,8 +636,9 @@ XLA_TEST_P(ReshapeTest, R4Dim0MinorLayoutToR2Dim0MajorLayout) {
LayoutUtil::MakeLayout({0, 1, 2, 3}));
// clang-format on
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Reshape(parameter, /*dimensions=*/{0, 1, 2, 3}, /*new_sizes=*/{2, 8});
@ -639,8 +672,9 @@ XLA_TEST_P(ReshapeTest, R2ToR4_3x8_To_3x2x1x4) {
{200, 201, 202, 203, 204, 205, 206, 207},
});
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Reshape(parameter, /*dimensions=*/{0, 1}, /*new_sizes=*/{3, 2, 1, 4});
// clang-format off
@ -666,8 +700,9 @@ XLA_TEST_P(ReshapeTest, R2ToR4_3x8_To_3x2x1x4_Dimensions_10) {
{200, 201, 202, 203, 204, 205, 206, 207},
});
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter));
Reshape(parameter, /*dimensions=*/{1, 0}, /*new_sizes=*/{3, 2, 1, 4});
// clang-format off
@ -694,8 +729,9 @@ XLA_TEST_P(ReshapeTest, R4ToR2_2x1x1x1_To_2x1) {
Literal input_literal = LiteralUtil::CreateR4FromArray4DWithLayout(
input, LayoutUtil::MakeLayout({3, 2, 1, 0}));
XlaOp parameter;
auto input_data = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(auto input_data,
CreateParameterAndTransferLiteral(
0, input_literal, "input", &builder, &parameter));
Reshape(parameter, /*dimensions=*/{0, 1, 2, 3}, /*new_sizes=*/{2, 1});
Literal expected = LiteralUtil::ReshapeSlice({2, 1}, {1, 0}, input_literal);
@ -713,8 +749,9 @@ XLA_TEST_P(ReshapeTest, R4ToR2_2x1x4x1_To_4x2) {
Literal input_literal = LiteralUtil::CreateR4FromArray4DWithLayout(
input, LayoutUtil::MakeLayout({3, 2, 1, 0}));
XlaOp parameter;
auto input_data = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(auto input_data,
CreateParameterAndTransferLiteral(
0, input_literal, "input", &builder, &parameter));
Reshape(parameter, /*dimensions=*/{0, 1, 2, 3}, /*new_sizes=*/{4, 2});
Literal expected = LiteralUtil::ReshapeSlice({4, 2}, {1, 0}, input_literal);
@ -733,8 +770,9 @@ XLA_TEST_P(ReshapeTest, R4ToR2_5x10x2x3_To_5x60_Dimensions_0213) {
Literal input_literal = LiteralUtil::CreateR4FromArray4DWithLayout(
input, LayoutUtil::MakeLayout({3, 2, 1, 0}));
XlaOp parameter;
auto input_data = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(auto input_data,
CreateParameterAndTransferLiteral(
0, input_literal, "input", &builder, &parameter));
Reshape(parameter, /*dimensions=*/{0, 2, 1, 3},
/*new_sizes=*/{5, 60});
@ -759,8 +797,9 @@ XLA_TEST_P(ReshapeTest, NoopReshape) {
Literal input_literal = LiteralUtil::CreateR4FromArray4DWithLayout(
input_array, LayoutUtil::MakeLayout({1, 2, 3, 0}));
XlaOp parameter;
auto input_data = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(auto input_data,
CreateParameterAndTransferLiteral(
0, input_literal, "input", &builder, &parameter));
Reshape(parameter, /*dimensions=*/{3, 0, 1, 2},
/*new_sizes=*/{7, 2, 3, 5});
XlaComputation computation = builder.Build().ConsumeValueOrDie();
@ -793,8 +832,9 @@ XLA_TEST_P(ReshapeTest, R4ToR4Reshape_Trivial) {
{{13, 14, 15, 16}, {17, 18, 19, 20}, {21, 22, 23, 24}}}});
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, literal_1x2x3x4, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, literal_1x2x3x4, "input",
&builder, &parameter));
Reshape(parameter, /*dimensions=*/{0, 1, 2, 3},
/*new_sizes=*/{1, 2, 3, 4});
@ -808,8 +848,9 @@ XLA_TEST_P(ReshapeTest, R4ToR4Reshape) {
XlaBuilder builder(TestName());
XlaOp parameter;
auto input = CreateParameterAndTransferLiteral(0, literal_1x2x3x4, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(
auto input, CreateParameterAndTransferLiteral(0, literal_1x2x3x4, "input",
&builder, &parameter));
Reshape(parameter, /*dimensions=*/{1, 3, 2, 0},
/*new_sizes=*/{2, 4, 3, 1});
@ -840,8 +881,9 @@ XLA_TEST_P(ReshapeTest, R4TwoMinorTransposeSimple) {
input, LayoutUtil::MakeLayout({3, 2, 1, 0}));
XlaBuilder builder(TestName());
XlaOp parameter;
auto input_data = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(auto input_data,
CreateParameterAndTransferLiteral(
0, input_literal, "input", &builder, &parameter));
Reshape(parameter, /*dimensions=*/{0, 1, 3, 2},
/*new_sizes=*/new_bounds);
@ -867,8 +909,9 @@ XLA_TEST_P(ReshapeTest, R4TwoMinorTransposeMajorFirstEffectiveR2) {
input, LayoutUtil::MakeLayout({3, 2, 1, 0}));
XlaBuilder builder(TestName());
XlaOp parameter;
auto input_data = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(auto input_data,
CreateParameterAndTransferLiteral(
0, input_literal, "input", &builder, &parameter));
Reshape(parameter, /*dimensions=*/{0, 1, 3, 2},
/*new_sizes=*/new_bounds);
@ -894,8 +937,9 @@ XLA_TEST_P(ReshapeTest, R4TwoMinorTransposeMajorFirstMinorEffectiveR1) {
input, LayoutUtil::MakeLayout({3, 2, 1, 0}));
XlaBuilder builder(TestName());
XlaOp parameter;
auto input_data = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(auto input_data,
CreateParameterAndTransferLiteral(
0, input_literal, "input", &builder, &parameter));
Reshape(parameter, /*dimensions=*/{0, 1, 3, 2},
/*new_sizes=*/new_bounds);
@ -922,8 +966,9 @@ XLA_TEST_P(ReshapeTest, R4TwoMinorTransposeMajorFirstMinorEffectiveR1InR2) {
input, LayoutUtil::MakeLayout({3, 2, 1, 0}));
XlaBuilder builder(TestName());
XlaOp parameter;
auto input_data = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(auto input_data,
CreateParameterAndTransferLiteral(
0, input_literal, "input", &builder, &parameter));
Reshape(parameter, /*dimensions=*/{0, 1, 3, 2},
/*new_sizes=*/new_bounds);
@ -949,8 +994,9 @@ XLA_TEST_P(ReshapeTest, R4TwoMinorTransposeTrivialR2) {
input, LayoutUtil::MakeLayout({0, 1, 2, 3}));
XlaBuilder builder(TestName());
XlaOp parameter;
auto input_data = CreateParameterAndTransferLiteral(0, input_literal, "input",
&builder, &parameter);
TF_ASSERT_OK_AND_ASSIGN(auto input_data,
CreateParameterAndTransferLiteral(
0, input_literal, "input", &builder, &parameter));
Reshape(parameter, /*dimensions=*/{1, 0, 2, 3},
/*new_sizes=*/new_bounds);