Function XLA_TEST_P

tensorflow/compiler/xla/tests/batch_normalization_test.cc:104–118

#endif

XLA_TEST_P(BatchNormalizationTest, SubtractInZ) {
  XlaBuilder builder("subtract_in_z_one_sample");
  auto x = ConstantLiteral(&builder, input_literal_);
  auto y = ConstantR1<float>(&builder, {3.14, 4.25});
  Sub(x, y, /*broadcast_dimensions=*/{1});

  Array4D<float> expected(kSamples, kZ, kY, kX);
  Array2D<float> pz({
      {-1.0f - 3.14f, 4.1f - 4.25f},  // p0
      {2.0f - 3.14f, 4.1f - 4.25f},   // p1
      {5.0f - 3.14f, 4.4f - 4.25f},   // p2
  });
  expected.FillWithPZ(pz);
  ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_);
}

XLA_TEST_P(BatchNormalizationTest, SquareTesseractElementwise) {
  XlaBuilder builder("square_tesseract_elementwise");