Function XLA_TEST_P

tensorflow/compiler/xla/tests/reshape_test.cc:54–67

Collapses 2-dimensional pseudo-scalar (single-element array) to 1 dimension.

// Collapses 2-dimensional pseudo-scalar (single-element array) to 1 dimension.
XLA_TEST_P(ReshapeTest, CollapseTrivial1x1) {
  XlaBuilder builder(TestName());
  Array2D<float> input_array(1, 1);
  input_array.Fill(1.0f);
  auto input_literal = LiteralUtil::CreateR2FromArray2D(input_array);
  XlaOp parameter;
  auto input = CreateParameterAndTransferLiteral(0, input_literal, "parameter",
                                                 &builder, &parameter);
  Collapse(/*operand=*/parameter, /*dimensions=*/{0, 1});

  auto expected_literal = LiteralUtil::CreateR1<float>({1.0f});
  ComputeAndCompareLiteral(&builder, expected_literal, {input.get()},
                           zero_error_spec_);
}

XLA_TEST_P(ReshapeTest, CollapseTrivialR1EmptyDims) {
  XlaBuilder builder(TestName());