<commit_before>/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

// Tests the reduce-window XLA operation.

#include <limits>
#include <memory>

#include "tensorflow/compiler/xla/array2d.h"
#include "tensorflow/compiler/xla/array3d.h"
#include "tensorflow/compiler/xla/array4d.h"
#include "tensorflow/compiler/xla/client/computation_builder.h"
#include "tensorflow/compiler/xla/client/lib/arithmetic.h"
#include "tensorflow/compiler/xla/client/local_client.h"
#include "tensorflow/compiler/xla/client/padding.h"
#include "tensorflow/compiler/xla/reference_util.h"
#include "tensorflow/compiler/xla/shape_util.h"
#include "tensorflow/compiler/xla/tests/client_library_test_base.h"
#include "tensorflow/compiler/xla/tests/hlo_test_base.h"
#include "tensorflow/compiler/xla/tests/literal_test_util.h"
#include "tensorflow/compiler/xla/tests/test_macros.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
#include "tensorflow/core/lib/gtl/array_slice.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/types.h"

namespace xla {
namespace {

class ReduceWindowTest : public ClientLibraryTestBase {
 public:
  ReduceWindowTest() : builder_(client_, TestName()) {}

  void ReduceWindowAdd(const ComputationDataHandle& input,
                       tensorflow::gtl::ArraySlice<int64> window_dimensions,
                       tensorflow::gtl::ArraySlice<int64> window_strides,
                       Padding padding) {
    builder_.ReduceWindow(input, builder_.ConstantR0<float>(0.0f),
                          CreateScalarAddComputation(F32, &builder_),
                          window_dimensions, window_strides, padding);
  }

  void ReduceWindowMax(const ComputationDataHandle& input,
                       tensorflow::gtl::ArraySlice<int64> window_dimensions,
                       tensorflow::gtl::ArraySlice<int64> window_strides,
                       Padding padding) {
    builder_.ReduceWindow(
        input, builder_.ConstantLiteral(Literal::MinValue(F32)),
        CreateScalarMax(), window_dimensions, window_strides, padding);
  }

  void ReduceWindowMin(const ComputationDataHandle& input,
                       tensorflow::gtl::ArraySlice<int64> window_dimensions,
                       tensorflow::gtl::ArraySlice<int64> window_strides,
                       Padding padding) {
    builder_.ReduceWindow(input,
                          builder_.ConstantLiteral(Literal::MaxValue(F32)),
                          CreateScalarMinComputation(F32, &builder_),
                          window_dimensions, window_strides, padding);
  }

  ComputationBuilder builder_;
};

TEST_F(ReduceWindowTest, Min3In5Stride2) {
  const auto input = builder_.ConstantR1<float>({10000, 1000, 100, 10, 1});
  ReduceWindowMin(input, {3}, {2}, Padding::kValid);
  ComputeAndCompareR1<float>(&builder_, {100, 1}, {}, ErrorSpec(0.0001));
}

XLA_TEST_F(ReduceWindowTest, ZeroElementSmall) {
  Array4D<float> input_array(1, 0, 2, 1);

  const auto input = builder_.ConstantR4FromArray4D<float>(input_array);
  Padding padding = Padding::kSame;
  ReduceWindowAdd(input, {1, 1, 2, 1}, {1, 1, 1, 1}, padding);

  auto res = ReferenceUtil::ReduceWindow4DAdd(input_array, 0.0f, {1, 1, 2, 1},
                                              {1, 1, 1, 1}, padding);

  ComputeAndCompareR4<float>(&builder_, *res, {}, ErrorSpec(1e-3, 1e-3));
}

TEST_F(ReduceWindowTest, NonSquareSmall) {
  Array4D<float> input_array(1, 2, 2, 1);
  input_array.FillRandom(2.f);

  const auto input = builder_.ConstantR4FromArray4D<float>(input_array);
  Padding padding = Padding::kSame;
  ReduceWindowAdd(input, {1, 1, 2, 1}, {1, 1, 1, 1}, padding);

  auto res = ReferenceUtil::ReduceWindow4DAdd(input_array, 0.0f, {1, 1, 2, 1},
                                              {1, 1, 1, 1}, padding);

  ComputeAndCompareR4<float>(&builder_, *res, {}, ErrorSpec(1e-3, 1e-3));
}

TEST_F(ReduceWindowTest, MiddleDimsSmall) {
  Array4D<float> input_array(1, 3, 3, 1);
  input_array.FillRandom(2.f);

  const auto input = builder_.ConstantR4FromArray4D<float>(input_array);
  Padding padding = Padding::kSame;
  ReduceWindowAdd(input, {1, 1, 1, 1}, {1, 2, 2, 1}, padding);

  auto res = ReferenceUtil::ReduceWindow4DAdd(input_array, 0.0f, {1, 1, 1, 1},
                                              {1, 2, 2, 1}, padding);

  ComputeAndCompareR4<float>(&builder_, *res, {}, ErrorSpec(1e-3, 1e-3));
}

TEST_F(ReduceWindowTest, Along2ndMinorDim) {
  Array4D<float> input_array(3, 6, 7, 32);
  input_array.FillRandom(2.f);

  // The parameters of this reduction mimic feature norm (e.g. LRN).
  int lrn_diameter = 7;  // diameter = 2*radius + 1 --> must be odd
  const auto input = builder_.ConstantR4FromArray4D<float>(input_array);
  Padding padding = Padding::kSame;
  ReduceWindowAdd(input, {1, 1, lrn_diameter, 1}, {1, 1, 1, 1}, padding);

  auto res = ReferenceUtil::ReduceWindow4DAdd(
      input_array, 0.0f, {1, 1, lrn_diameter, 1}, {1, 1, 1, 1}, padding);

  ComputeAndCompareR4<float>(&builder_, *res, {}, ErrorSpec(1e-3, 1e-3));
}

TEST_F(ReduceWindowTest, AmongMajor2Dims) {
  Array4D<float> input_array(4, 4, 6, 8);
  input_array.FillWithMinorDimNum();

  int win_len = 3;
  int win_stride = 1;

  Padding padding = Padding::kSame;
  const auto input_data_handle =
      builder_.ConstantR4FromArray4D<float>(input_array);
  // Reduce only along the x and y dimensions, according to the win_len.
  ReduceWindowAdd(input_data_handle, {win_len, win_len, 1, 1},
                  {win_stride, win_stride, 1, 1}, padding);

  auto result = ReferenceUtil::ReduceWindow4DAdd(
      input_array, 0.0f, {win_len, win_len, 1, 1},
      {win_stride, win_stride, 1, 1}, padding);
  ComputeAndCompareR4<float>(&builder_, *result, {}, ErrorSpec(1e-3, 1e-3));
}

TEST_F(ReduceWindowTest, AmongMajor2DimsMediumSize) {
  Array4D<float> input_array(9, 12, 4, 89);
  input_array.FillRandom(2.0f);

  int win_len = 3;
  int win_stride = 2;

  const auto input_data_handle =
      builder_.ConstantR4FromArray4D<float>(input_array);

  Padding padding = Padding::kSame;
  // Reduce only along the x and y dimensions, according to the win_len.
  ReduceWindowAdd(input_data_handle, {win_len, win_len, 1, 1},
                  {win_stride, win_stride, 1, 1}, padding);

  auto result = ReferenceUtil::ReduceWindow4DAdd(
      input_array, 0.0f, {win_len, win_len, 1, 1},
      {win_stride, win_stride, 1, 1}, padding);

  ComputeAndCompareR4<float>(&builder_, *result, {}, ErrorSpec(1e-3, 1e-3));
}

// TODO(b/32173947): Test support for arbitrary-sized padding.
TEST_F(ReduceWindowTest, DISABLED_AmongMajor2DimsMediumSizeLargePadding) {
  Array4D<float> input_array(9, 12, 4, 89);  // simulate Dim0IsMinor layout
  input_array.FillRandom(2.0f);

  int64 rank = 4;
  int win_len = 3;
  int win_stride = 2;

  const auto input_data_handle =
      builder_.ConstantR4FromArray4D<float>(input_array);

  Padding padding = Padding::kSame;
  // Reduce only along the x and y dimensions, according to the win_len.
  // Create padding vector with large padding values in the reduction dims.
  std::vector<std::pair<int64, int64>> low_high_padding;
  low_high_padding.resize(rank, {4, 4});

  builder_.ReduceWindowWithGeneralPadding(
      input_data_handle, builder_.ConstantR0<float>(0.0f),
      CreateScalarAddComputation(F32, &builder_), {win_len, win_len, 1, 1},
      {win_stride, win_stride, 1, 1}, low_high_padding);

  auto result = ReferenceUtil::ReduceWindow4DAdd(
      input_array, 0.0f, {win_len, win_len, 1, 1},
      {win_stride, win_stride, 1, 1}, padding);

  ComputeAndCompareR4<float>(&builder_, *result, {}, ErrorSpec(1e-3, 1e-3));
}

XLA_TEST_F(ReduceWindowTest, Add1x1x2In2x1x2) {
  Array3D<float> input_array(2, 1, 2);
  input_array(0, 0, 0) = 1000;
  input_array(0, 0, 1) = 100;
  input_array(1, 0, 0) = 10;
  input_array(1, 0, 1) = 1;
  auto input = builder_.ConstantR3FromArray3D<float>(input_array);

  ReduceWindowAdd(input, {1, 1, 2}, {1, 1, 1}, Padding::kValid);

  Array3D<float> expected(2, 1, 1);
  expected(0, 0, 0) = 1100;
  expected(1, 0, 0) = 11;
  ComputeAndCompareR3<float>(&builder_, expected, {}, ErrorSpec(0.0001));
}

XLA_TEST_F(ReduceWindowTest, Add1x1x2In2x1x3Stride1x1x2) {
  Array3D<float> input_array(2, 1, 3);
  input_array(0, 0, 0) = 100;
  input_array(0, 0, 1) = 10;
  input_array(0, 0, 2) = 1;
  input_array(1, 0, 0) = 500;
  input_array(1, 0, 1) = 50;
  input_array(1, 0, 2) = 5;
  auto input = builder_.ConstantR3FromArray3D<float>(input_array);

  ReduceWindowAdd(input, {1, 1, 2}, {1, 1, 2}, Padding::kValid);

  Array3D<float> expected(2, 1, 1);
  expected(0, 0, 0) = 110;
  expected(1, 0, 0) = 550;
  ComputeAndCompareR3<float>(&builder_, expected, {}, ErrorSpec(0.0001));
}

XLA_TEST_F(ReduceWindowTest, Add1x1x2In2x1x3SamePad) {
  Array3D<float> input_array(2, 1, 3);
  input_array(0, 0, 0) = 100;
  input_array(0, 0, 1) = 10;
  input_array(0, 0, 2) = 1;
  input_array(1, 0, 0) = 500;
  input_array(1, 0, 1) = 50;
  input_array(1, 0, 2) = 5;
  auto input = builder_.ConstantR3FromArray3D<float>(input_array);

  ReduceWindowAdd(input, {1, 1, 2}, {1, 1, 1}, Padding::kSame);

  Array3D<float> expected(2, 1, 3);
  expected(0, 0, 0) = 110;
  expected(0, 0, 1) = 11;
  expected(0, 0, 2) = 1;
  expected(1, 0, 0) = 550;
  expected(1, 0, 1) = 55;
  expected(1, 0, 2) = 5;
  ComputeAndCompareR3<float>(&builder_, expected, {}, ErrorSpec(0.0001));
}

// Tests a reduction function that is not a simple add/min/max/etc.
XLA_TEST_F(ReduceWindowTest, NonstandardReduceFunction) {
  Array4D<float> input_array(1, 2, 2, 1);
  input_array(0, 0, 0, 0) = 1;
  input_array(0, 0, 1, 0) = 2;
  input_array(0, 1, 0, 0) = 3;
  input_array(0, 1, 1, 0) = 4;

  const auto input = builder_.ConstantR4FromArray4D<float>(input_array);
  Padding padding = Padding::kValid;

  const Shape scalar = ShapeUtil::MakeShape(F32, {});
  auto b = builder_.CreateSubBuilder("unusual");
  auto lhs = b->Parameter(0, scalar, "lhs");
  auto rhs = b->Parameter(1, scalar, "rhs");
  b->Min(b->Add(lhs, rhs), b->ConstantR0<float>(8.0f));
  Computation reduce_fn = b->BuildAndNoteError();

  builder_.ReduceWindow(input, builder_.ConstantR0<float>(3.0f), reduce_fn,
                        /*window_dimensions=*/{1, 1, 2, 1},
                        /*window_strides=*/{1, 1, 1, 1}, padding);

  const auto reduce_func = [](float arg1, float arg2) {
    return std::min<float>(arg1 + arg2, 8.0f);
  };

  auto expected =
      ReferenceUtil::ReduceWindow4DGeneric(input_array, 3.0f, reduce_func,
                                           /*window=*/{1, 1, 2, 1},
                                           /*stride=*/{1, 1, 1, 1}, padding);

  ComputeAndCompareR4<float>(&builder_, *expected, {}, ErrorSpec(1e-3, 1e-3));
}

TEST_F(ReduceWindowTest, R4UnitWindow) {
  Array4D<float> input_array(13, 12, 8, 15);
  input_array.Fill(1.0f);
  std::unique_ptr<Literal> input_literal =
      Literal::CreateR4FromArray4DWithLayout(
          input_array, LayoutUtil::MakeLayout({0, 3, 2, 1}));
  ComputationDataHandle input =
      builder_.Parameter(0, input_literal->shape(), "operand");

  Padding padding = Padding::kSame;
  ReduceWindowAdd(input, {1, 1, 7, 1}, {1, 4, 1, 1}, padding);

  auto res = ReferenceUtil::ReduceWindow4DAdd(input_array, 0.0f, {1, 1, 7, 1},
                                              {1, 4, 1, 1}, padding);

  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<GlobalData> input_data,
                          client_->TransferToServer(*input_literal));
  ComputeAndCompareR4<float>(&builder_, *res, {input_data.get()},
                             ErrorSpec(1e-3, 1e-3));
}

XLA_TEST_F(HloTestBase, R6Add) {
  auto b = HloComputation::Builder(TestName());

  std::vector<int64> input_dims(6, 8);
  std::unique_ptr<Literal> arg_literal =
      Literal::CreateFullWithMonotonicDim0MajorLayout<float>(input_dims, 1.0f);
  auto input =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(arg_literal)));

  auto init_value = b.AddInstruction(
      HloInstruction::CreateConstant(Literal::CreateR0<float>(0.f)));

  HloComputation::Builder add_computation("add");
  Shape scalar_shape = ShapeUtil::MakeShape(F32, {});
  auto param_lhs = add_computation.AddInstruction(
      HloInstruction::CreateParameter(0, scalar_shape, "lhs"));
  auto param_rhs = add_computation.AddInstruction(
      HloInstruction::CreateParameter(1, scalar_shape, "rhs"));
  add_computation.AddInstruction(HloInstruction::CreateBinary(
      scalar_shape, HloOpcode::kAdd, param_lhs, param_rhs));<commit_msg>Adds test for six-dimensional pooling case.<commit_after>/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

// Tests the reduce-window XLA operation.

#include <limits>
#include <memory>

#include "tensorflow/compiler/xla/array2d.h"
#include "tensorflow/compiler/xla/array3d.h"
#include "tensorflow/compiler/xla/array4d.h"
#include "tensorflow/compiler/xla/client/computation_builder.h"
#include "tensorflow/compiler/xla/client/lib/arithmetic.h"
#include "tensorflow/compiler/xla/client/local_client.h"
#include "tensorflow/compiler/xla/client/padding.h"
#include "tensorflow/compiler/xla/reference_util.h"
#include "tensorflow/compiler/xla/shape_util.h"
#include "tensorflow/compiler/xla/tests/client_library_test_base.h"
#include "tensorflow/compiler/xla/tests/hlo_test_base.h"
#include "tensorflow/compiler/xla/tests/literal_test_util.h"
#include "tensorflow/compiler/xla/tests/test_macros.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/gtl/array_slice.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/types.h"

namespace xla {
namespace {

class ReduceWindowTest : public ClientLibraryTestBase {
 public:
  ReduceWindowTest() : builder_(client_, TestName()) {}

  void ReduceWindowAdd(const ComputationDataHandle& input,
                       tensorflow::gtl::ArraySlice<int64> window_dimensions,
                       tensorflow::gtl::ArraySlice<int64> window_strides,
                       Padding padding) {
    builder_.ReduceWindow(input, builder_.ConstantR0<float>(0.0f),
                          CreateScalarAddComputation(F32, &builder_),
                          window_dimensions, window_strides, padding);
  }

  void ReduceWindowMax(const ComputationDataHandle& input,
                       tensorflow::gtl::ArraySlice<int64> window_dimensions,
                       tensorflow::gtl::ArraySlice<int64> window_strides,
                       Padding padding) {
    builder_.ReduceWindow(
        input, builder_.ConstantLiteral(Literal::MinValue(F32)),
        CreateScalarMax(), window_dimensions, window_strides, padding);
  }

  void ReduceWindowMin(const ComputationDataHandle& input,
                       tensorflow::gtl::ArraySlice<int64> window_dimensions,
                       tensorflow::gtl::ArraySlice<int64> window_strides,
                       Padding padding) {
    builder_.ReduceWindow(input,
                          builder_.ConstantLiteral(Literal::MaxValue(F32)),
                          CreateScalarMinComputation(F32, &builder_),
                          window_dimensions, window_strides, padding);
  }

  ComputationBuilder builder_;
};

TEST_F(ReduceWindowTest, Min3In5Stride2) {
  const auto input = builder_.ConstantR1<float>({10000, 1000, 100, 10, 1});
  ReduceWindowMin(input, {3}, {2}, Padding::kValid);
  ComputeAndCompareR1<float>(&builder_, {100, 1}, {}, ErrorSpec(0.0001));
}

XLA_TEST_F(ReduceWindowTest, ZeroElementSmall) {
  Array4D<float> input_array(1, 0, 2, 1);

  const auto input = builder_.ConstantR4FromArray4D<float>(input_array);
  Padding padding = Padding::kSame;
  ReduceWindowAdd(input, {1, 1, 2, 1}, {1, 1, 1, 1}, padding);

  auto res = ReferenceUtil::ReduceWindow4DAdd(input_array, 0.0f, {1, 1, 2, 1},
                                              {1, 1, 1, 1}, padding);

  ComputeAndCompareR4<float>(&builder_, *res, {}, ErrorSpec(1e-3, 1e-3));
}

TEST_F(ReduceWindowTest, NonSquareSmall) {
  Array4D<float> input_array(1, 2, 2, 1);
  input_array.FillRandom(2.f);

  const auto input = builder_.ConstantR4FromArray4D<float>(input_array);
  Padding padding = Padding::kSame;
  ReduceWindowAdd(input, {1, 1, 2, 1}, {1, 1, 1, 1}, padding);

  auto res = ReferenceUtil::ReduceWindow4DAdd(input_array, 0.0f, {1, 1, 2, 1},
                                              {1, 1, 1, 1}, padding);

  ComputeAndCompareR4<float>(&builder_, *res, {}, ErrorSpec(1e-3, 1e-3));
}

TEST_F(ReduceWindowTest, MiddleDimsSmall) {
  Array4D<float> input_array(1, 3, 3, 1);
  input_array.FillRandom(2.f);

  const auto input = builder_.ConstantR4FromArray4D<float>(input_array);
  Padding padding = Padding::kSame;
  ReduceWindowAdd(input, {1, 1, 1, 1}, {1, 2, 2, 1}, padding);

  auto res = ReferenceUtil::ReduceWindow4DAdd(input_array, 0.0f, {1, 1, 1, 1},
                                              {1, 2, 2, 1}, padding);

  ComputeAndCompareR4<float>(&builder_, *res, {}, ErrorSpec(1e-3, 1e-3));
}

TEST_F(ReduceWindowTest, Along2ndMinorDim) {
  Array4D<float> input_array(3, 6, 7, 32);
  input_array.FillRandom(2.f);

  // The parameters of this reduction mimic feature norm (e.g. LRN).
  int lrn_diameter = 7;  // diameter = 2*radius + 1 --> must be odd
  const auto input = builder_.ConstantR4FromArray4D<float>(input_array);
  Padding padding = Padding::kSame;
  ReduceWindowAdd(input, {1, 1, lrn_diameter, 1}, {1, 1, 1, 1}, padding);

  auto res = ReferenceUtil::ReduceWindow4DAdd(
      input_array, 0.0f, {1, 1, lrn_diameter, 1}, {1, 1, 1, 1}, padding);

  ComputeAndCompareR4<float>(&builder_, *res, {}, ErrorSpec(1e-3, 1e-3));
}

TEST_F(ReduceWindowTest, AmongMajor2Dims) {
  Array4D<float> input_array(4, 4, 6, 8);
  input_array.FillWithMinorDimNum();

  int win_len = 3;
  int win_stride = 1;

  Padding padding = Padding::kSame;
  const auto input_data_handle =
      builder_.ConstantR4FromArray4D<float>(input_array);
  // Reduce only along the x and y dimensions, according to the win_len.
  ReduceWindowAdd(input_data_handle, {win_len, win_len, 1, 1},
                  {win_stride, win_stride, 1, 1}, padding);

  auto result = ReferenceUtil::ReduceWindow4DAdd(
      input_array, 0.0f, {win_len, win_len, 1, 1},
      {win_stride, win_stride, 1, 1}, padding);
  ComputeAndCompareR4<float>(&builder_, *result, {}, ErrorSpec(1e-3, 1e-3));
}

TEST_F(ReduceWindowTest, AmongMajor2DimsMediumSize) {
  Array4D<float> input_array(9, 12, 4, 89);
  input_array.FillRandom(2.0f);

  int win_len = 3;
  int win_stride = 2;

  const auto input_data_handle =
      builder_.ConstantR4FromArray4D<float>(input_array);

  Padding padding = Padding::kSame;
  // Reduce only along the x and y dimensions, according to the win_len.
  ReduceWindowAdd(input_data_handle, {win_len, win_len, 1, 1},
                  {win_stride, win_stride, 1, 1}, padding);

  auto result = ReferenceUtil::ReduceWindow4DAdd(
      input_array, 0.0f, {win_len, win_len, 1, 1},
      {win_stride, win_stride, 1, 1}, padding);

  ComputeAndCompareR4<float>(&builder_, *result, {}, ErrorSpec(1e-3, 1e-3));
}

// TODO(b/32173947): Test support for arbitrary-sized padding.
TEST_F(ReduceWindowTest, DISABLED_AmongMajor2DimsMediumSizeLargePadding) {
  Array4D<float> input_array(9, 12, 4, 89);  // simulate Dim0IsMinor layout
  input_array.FillRandom(2.0f);

  int64 rank = 4;
  int win_len = 3;
  int win_stride = 2;

  const auto input_data_handle =
      builder_.ConstantR4FromArray4D<float>(input_array);

  Padding padding = Padding::kSame;
  // Reduce only along the x and y dimensions, according to the win_len.
  // Create padding vector with large padding values in the reduction dims.
  std::vector<std::pair<int64, int64>> low_high_padding;
  low_high_padding.resize(rank, {4, 4});

  builder_.ReduceWindowWithGeneralPadding(
      input_data_handle, builder_.ConstantR0<float>(0.0f),
      CreateScalarAddComputation(F32, &builder_), {win_len, win_len, 1, 1},
      {win_stride, win_stride, 1, 1}, low_high_padding);

  auto result = ReferenceUtil::ReduceWindow4DAdd(
      input_array, 0.0f, {win_len, win_len, 1, 1},
      {win_stride, win_stride, 1, 1}, padding);

  ComputeAndCompareR4<float>(&builder_, *result, {}, ErrorSpec(1e-3, 1e-3));
}

XLA_TEST_F(ReduceWindowTest, Add1x1x2In2x1x2) {
  Array3D<float> input_array(2, 1, 2);
  input_array(0, 0, 0) = 1000;
  input_array(0, 0, 1) = 100;
  input_array(1, 0, 0) = 10;
  input_array(1, 0, 1) = 1;
  auto input = builder_.ConstantR3FromArray3D<float>(input_array);

  ReduceWindowAdd(input, {1, 1, 2}, {1, 1, 1}, Padding::kValid);

  Array3D<float> expected(2, 1, 1);
  expected(0, 0, 0) = 1100;
  expected(1, 0, 0) = 11;
  ComputeAndCompareR3<float>(&builder_, expected, {}, ErrorSpec(0.0001));
}

XLA_TEST_F(ReduceWindowTest, Add1x1x2In2x1x3Stride1x1x2) {
  Array3D<float> input_array(2, 1, 3);
  input_array(0, 0, 0) = 100;
  input_array(0, 0, 1) = 10;
  input_array(0, 0, 2) = 1;
  input_array(1, 0, 0) = 500;
  input_array(1, 0, 1) = 50;
  input_array(1, 0, 2) = 5;
  auto input = builder_.ConstantR3FromArray3D<float>(input_array);

  ReduceWindowAdd(input, {1, 1, 2}, {1, 1, 2}, Padding::kValid);

  Array3D<float> expected(2, 1, 1);
  expected(0, 0, 0) = 110;
  expected(1, 0, 0) = 550;
  ComputeAndCompareR3<float>(&builder_, expected, {}, ErrorSpec(0.0001));
}

XLA_TEST_F(ReduceWindowTest, Add1x1x2In2x1x3SamePad) {
  Array3D<float> input_array(2, 1, 3);
  input_array(0, 0, 0) = 100;
  input_array(0, 0, 1) = 10;
  input_array(0, 0, 2) = 1;
  input_array(1, 0, 0) = 500;
  input_array(1, 0, 1) = 50;
  input_array(1, 0, 2) = 5;
  auto input = builder_.ConstantR3FromArray3D<float>(input_array);

  ReduceWindowAdd(input, {1, 1, 2}, {1, 1, 1}, Padding::kSame);

  Array3D<float> expected(2, 1, 3);
  expected(0, 0, 0) = 110;
  expected(0, 0, 1) = 11;
  expected(0, 0, 2) = 1;
  expected(1, 0, 0) = 550;
  expected(1, 0, 1) = 55;
  expected(1, 0, 2) = 5;
  ComputeAndCompareR3<float>(&builder_, expected, {}, ErrorSpec(0.0001));
}

// Tests a reduction function that is not a simple add/min/max/etc.
XLA_TEST_F(ReduceWindowTest, NonstandardReduceFunction) {
  Array4D<float> input_array(1, 2, 2, 1);
  input_array(0, 0, 0, 0) = 1;
  input_array(0, 0, 1, 0) = 2;
  input_array(0, 1, 0, 0) = 3;
  input_array(0, 1, 1, 0) = 4;

  const auto input = builder_.ConstantR4FromArray4D<float>(input_array);
  Padding padding = Padding::kValid;

  const Shape scalar = ShapeUtil::MakeShape(F32, {});
  auto b = builder_.CreateSubBuilder("unusual");
  auto lhs = b->Parameter(0, scalar, "lhs");
  auto rhs = b->Parameter(1, scalar, "rhs");
  b->Min(b->Add(lhs, rhs), b->ConstantR0<float>(8.0f));
  Computation reduce_fn = b->BuildAndNoteError();

  builder_.ReduceWindow(input, builder_.ConstantR0<float>(3.0f), reduce_fn,
                        /*window_dimensions=*/{1, 1, 2, 1},
                        /*window_strides=*/{1, 1, 1, 1}, padding);

  const auto reduce_func = [](float arg1, float arg2) {
    return std::min<float>(arg1 + arg2, 8.0f);
  };

  auto expected =
      ReferenceUtil::ReduceWindow4DGeneric(input_array, 3.0f, reduce_func,
                                           /*window=*/{1, 1, 2, 1},
                                           /*stride=*/{1, 1, 1, 1}, padding);

  ComputeAndCompareR4<float>(&builder_, *expected, {}, ErrorSpec(1e-3, 1e-3));
}

TEST_F(ReduceWindowTest, R4UnitWindow) {
  Array4D<float> input_array(13, 12, 8, 15);
  input_array.Fill(1.0f);
  std::unique_ptr<Literal> input_literal =
      Literal::CreateR4FromArray4DWithLayout(
          input_array, LayoutUtil::MakeLayout({0, 3, 2, 1}));
  ComputationDataHandle input =
      builder_.Parameter(0, input_literal->shape(), "operand");

  Padding padding = Padding::kSame;
  ReduceWindowAdd(input, {1, 1, 7, 1}, {1, 4, 1, 1}, padding);

  auto res = ReferenceUtil::ReduceWindow4DAdd(input_array, 0.0f, {1, 1, 7, 1},
                                              {1, 4, 1, 1}, padding);

  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<GlobalData> input_data,
                          client_->TransferToServer(*input_literal));
  ComputeAndCompareR4<float>(&builder_, *res, {input_data.get()},
                             ErrorSpec(1e-3, 1e-3));
}

XLA_TEST_F(HloTestBase, R6AddMultipleStrides) {
  auto b = HloComputation::Builder(TestName());

  std::vector<int64> input_dims(6, 8);
  auto shape = ShapeUtil::MakeShape(F32, input_dims);

  std::unique_ptr<Literal> arg_literal = Literal::CreateFromShape(shape);
  auto generator = [&](tensorflow::gtl::ArraySlice<int64> indexes) -> float {
    return 1.0f;
  };
  TF_EXPECT_OK(arg_literal->Populate<float>(generator));

  auto input =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(arg_literal)));

  auto init_value = b.AddInstruction(
      HloInstruction::CreateConstant(Literal::CreateR0<float>(0.f)));

  HloComputation::Builder add_computation("add");
  Shape scalar_shape = ShapeUtil::MakeShape(F32, {});
  auto param_lhs = add_computation.AddInstruction(
      HloInstruction::CreateParameter(0, scalar_shape, "lhs"));
  auto param_rhs = add_computation.AddInstruction(
      HloInstruction::CreateParameter(1, scalar_shape, "rhs"));
  add_computation.AddInstruction(HloInstruction::CreateBinary(
      scalar_shape, HloOpcode::kAdd, param_lhs, param_rhs));

  auto module = CreateNewModule();
  auto add_func = module->AddEmbeddedComputation(add_computation.Build());

  WindowDimension trivial_dim;
  trivial_dim.set_size(1);
  trivial_dim.set_stride(1);
  trivial_dim.set_padding_low(0);
  trivial_dim.set_padding_high(0);
  trivial_dim.set_window_dilation(1);
  trivial_dim.set_base_dilation(1);

  WindowDimension active_dim;
  active_dim.set_size(3);
  active_dim.set_stride(1);
  active_dim.set_padding_low(0);
  active_dim.set_padding_high(0);
  active_dim.set_window_dilation(1);
  active_dim.set_base_dilation(1);

  Window window;
  *window.add_dimensions() = active_dim;
  *window.add_dimensions() = trivial_dim;
  *window.add_dimensions() = active_dim;
  *window.add_dimensions() = active_dim;
  *window.add_dimensions() = trivial_dim;
  *window.add_dimensions() = trivial_dim;

  // Non-monotonic output layout with minor dims trivial.
  std::vector<int64> output_layout = {1, 5, 3, 2, 0, 4};
  std::vector<int64> output_dims = {6, 8, 6, 6, 8, 8};
  Shape result_shape =
      ShapeUtil::MakeShapeWithLayout(F32, output_dims, output_layout);
  b.AddInstruction(HloInstruction::CreateReduceWindow(
      result_shape, input, init_value, window, add_func));

  std::unique_ptr<Literal> expected = Literal::CreateFromShape(result_shape);
  auto out_generator =
      [&](tensorflow::gtl::ArraySlice<int64> indexes) -> float {
    return 27.0f;
  };
  TF_EXPECT_OK(expected->Populate<float>(out_generator));

  module->AddEntryComputation(b.Build());
  auto actual = ExecuteAndTransfer(std::move(module), {});

  LiteralTestUtil::ExpectNear(*actual, *expected, ErrorSpec(1e-3, 1e-3));
}

XLA_TEST_F(HloTestBase, R6Add) {
  auto b = HloComputation::Builder(TestName());

  std::vector<int64> input_dims(6, 8);
  std::unique_ptr<Literal> arg_literal =
      Literal::CreateFullWithMonotonicDim0MajorLayout<float>(input_dims, 1.0f);
  auto input =
      b.AddInstruction(HloInstruction::CreateConstant(std::move(arg_literal)));

  auto init_value = b.AddInstruction(
      HloInstruction::CreateConstant(Literal::CreateR0<float>(0.f)));

  HloComputation::Builder add_computation("add");
  Shape scalar_shape = ShapeUtil::MakeShape(F32, {});
  auto param_lhs = add_computation.AddInstruction(
      HloInstruction::CreateParameter(0, scalar_shape, "lhs"));
  auto param_rhs = add_computation.AddInstruction(
      HloInstruction::CreateParameter(1, scalar_shape, "rhs"));
  add_computation.AddInstruction(HloInstruction::CreateBinary(
      scalar_shape, HloOpcode::kAdd, param_lhs, param_rhs));
{'fraction_non_alphanumeric': 0.14278350515463917, 'fraction_numerical': 0.046597938144329894, 'mean_word_length': 4.046999653139091, 'pattern_counts': {'":': 0, '<': 135, '<?xml version=': 0, '>': 159, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 24, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3949777', 'n_tokens_mistral': 11076, 'n_tokens_neox': 10152, 'n_words': 2139}