<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.
==============================================================================*/

#include <memory>
#include <vector>

#include "tensorflow/compiler/xla/client/global_data.h"
#include "tensorflow/compiler/xla/client/local_client.h"
#include "tensorflow/compiler/xla/client/xla_builder.h"
#include "tensorflow/compiler/xla/tests/client_library_test_base.h"
#include "tensorflow/compiler/xla/tests/literal_test_util.h"
#include "tensorflow/compiler/xla/tests/test_macros.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/types.h"

namespace xla {
namespace {

class SelectTest : public ClientLibraryTestBase {
 public:
  ErrorSpec error_spec_{0.0001};
};

TEST_F(SelectTest, SelectScalarF32True) {
  XlaBuilder builder(TestName());
  auto pred = ConstantR0<bool>(&builder, true);
  auto on_true = ConstantR0<float>(&builder, 123.0f);
  auto on_false = ConstantR0<float>(&builder, 42.0f);
  Select(pred, on_true, on_false);

  ComputeAndCompareR0<float>(&builder, 123.0f, {}, error_spec_);
}

TEST_F(SelectTest, SelectScalarS32True) {
  XlaBuilder builder(TestName());
  auto pred = ConstantR0<bool>(&builder, true);
  auto on_true = ConstantR0<int32>(&builder, -42);
  auto on_false = ConstantR0<int32>(&builder, 42);
  Select(pred, on_true, on_false);

  ComputeAndCompareR0<int32>(&builder, -42, {});
}

TEST_F(SelectTest, SelectScalarF32False) {
  XlaBuilder builder(TestName());
  auto pred = ConstantR0<bool>(&builder, false);
  auto on_true = ConstantR0<float>(&builder, 123.0f);
  auto on_false = ConstantR0<float>(&builder, 42.0f);
  Select(pred, on_true, on_false);

  ComputeAndCompareR0<float>(&builder, 42.0f, {}, error_spec_);
}

XLA_TEST_F(SelectTest, SelectR1S0F32WithConstantR1S0PRED) {
  XlaBuilder builder(TestName());
  auto pred = ConstantR1<bool>(&builder, {});
  auto on_true = ConstantR1<float>(&builder, {});
  auto on_false = ConstantR1<float>(&builder, {});
  Select(pred, on_true, on_false);

  ComputeAndCompareR1<float>(&builder, {}, {}, error_spec_);
}

TEST_F(SelectTest, SelectR1F32WithConstantR1PRED) {
  XlaBuilder builder(TestName());
  auto pred = ConstantR1<bool>(&builder, {false, true, false, true, false});
  auto on_true =
      ConstantR1<float>(&builder, {-2.5f, 25.5f, 2.25f, -10.0f, 6.0f});
  auto on_false =
      ConstantR1<float>(&builder, {10.0f, 5.0f, 1.0f, 10.0f, -6.0f});
  Select(pred, on_true, on_false);

  ComputeAndCompareR1<float>(&builder, {10.0f, 25.5f, 1.0f, -10.0f, -6.0f}, {},
                             error_spec_);
}

XLA_TEST_F(SelectTest, SelectR1S0F32WithCmpR1S0S32s) {
  // Similar to SelectR1S0F32WithConstantR1S0PRED, except that the pred vector
  // is not a constant, but rather the result of comparing two other vectors.
  XlaBuilder builder(TestName());
  auto v1 = ConstantR1<int32>(&builder, {});
  auto v2 = ConstantR1<int32>(&builder, {});
  auto cmp = Eq(v1, v2);
  auto on_true = ConstantR1<float>(&builder, {});
  auto on_false = ConstantR1<float>(&builder, {});
  Select(cmp, on_true, on_false);

  ComputeAndCompareR1<float>(&builder, {}, {}, error_spec_);
}

TEST_F(SelectTest, SelectR1F32WithCmpR1S32s) {
  // Similar to SelectR1F32WithConstantR1PRED, except that the pred vector is
  // not a constant, but rather the result of comparing two other vectors.
  XlaBuilder builder(TestName());
  auto v1 = ConstantR1<int32>(&builder, {1, 2, 3, 4, 5});
  auto v2 = ConstantR1<int32>(&builder, {9, 2, 9, 4, 9});
  auto cmp = Eq(v1, v2);
  auto on_true =
      ConstantR1<float>(&builder, {-2.5f, 25.5f, 2.25f, -10.0f, 6.0f});
  auto on_false =
      ConstantR1<float>(&builder, {10.0f, 5.0f, 1.0f, 10.0f, -6.0f});
  Select(cmp, on_true, on_false);

  ComputeAndCompareR1<float>(&builder, {10.0f, 25.5f, 1.0f, -10.0f, -6.0f}, {},
                             error_spec_);
}

TEST_F(SelectTest, SelectR1F32WithCmpR1F32s) {
  // Similar to SelectR1F32WithCmpR1S32s, except "gt"-comparing two R1F32s.
  XlaBuilder builder(TestName());
  auto v1 = ConstantR1<float>(&builder, {1.0f, 2.0f, 3.0f, 4.0f, 5.0f});
  auto v2 = ConstantR1<float>(&builder, {-1.0f, -2.0f, 13.0f, 14.0f, 4.4f});
  auto cmp = Gt(v1, v2);
  auto on_true =
      ConstantR1<float>(&builder, {-2.5f, 25.5f, 2.25f, -10.0f, 6.0f});
  auto on_false =
      ConstantR1<float>(&builder, {10.0f, 5.0f, 1.0f, 10.0f, -6.0f});
  Select(cmp, on_true, on_false);

  ComputeAndCompareR1<float>(&builder, {-2.5f, 25.5f, 1.0f, 10.0f, 6.0f}, {},
                             error_spec_);
}

TEST_F(SelectTest, SelectR1F32WithCmpR1F32sFromParamsSmall) {
  // Selects among two R1F32s, which come from parameters. v1 and v2 are
  // compared, and selection between them happens based on a gt-comparison mask.
  XlaBuilder builder(TestName());

  XlaOp v1, v2;
  std::unique_ptr<GlobalData> param0_data = CreateR1Parameter<float>(
      {41.0f, 2.0f, 3.0f, 84.0f}, /*parameter_number=*/0, /*name=*/"v1",
      /*builder=*/&builder, /*data_handle=*/&v1);
  std::unique_ptr<GlobalData> param1_data = CreateR1Parameter<float>(
      {21.0f, 22.0f, 23.0f, 24.0f}, /*parameter_number=*/1, /*name=*/"v2",
      /*builder=*/&builder, /*data_handle=*/&v2);

  auto cmp = Gt(v1, v2);
  Select(cmp, v1, v2);
  ComputeAndCompareR1<float>(&builder, {41.0f, 22.0f, 23.0f, 84.0f},
                             {param0_data.get(), param1_data.get()},
                             error_spec_);
}

TEST_F(SelectTest, SelectR1F32WithCmpR1F32sFromParamsLarge) {
  // Similar to SelectR1F32WithCmpR1F32sFromParamsSmall, except that the
  // data size passed in and out is large.
  XlaBuilder builder(TestName());

  // Number of floats in the data passed into and out of the computation.
  constexpr int datalen = 15 * 1000;

  // The inputs are initialized with a special pattern where in the first third
  // of the data v1[i] > v2[i] and elsewhere it's vice versa.
  std::vector<float> v1vec;
  std::vector<float> v2vec;
  std::vector<float> expected_vec;
  for (int i = 0; i < datalen; ++i) {
    float smaller = i;
    float larger = i * 2;
    if (i < datalen / 3) {
      v1vec.push_back(larger);
      v2vec.push_back(smaller);
    } else {
      v1vec.push_back(smaller);
      v2vec.push_back(larger);
    }
    expected_vec.push_back(larger);
  }

  XlaOp v1, v2;
  std::unique_ptr<GlobalData> param0_data =
      CreateR1Parameter<float>(v1vec, /*parameter_number=*/0, /*name=*/"v1",
                               /*builder=*/&builder, /*data_handle=*/&v1);
  std::unique_ptr<GlobalData> param1_data =
      CreateR1Parameter<float>(v2vec, /*parameter_number=*/1, /*name=*/"v2",
                               /*builder=*/&builder, /*data_handle=*/&v2);

  auto cmp = Gt(v1, v2);
  Select(cmp, v1, v2);
  ComputeAndCompareR1<float>(&builder, expected_vec,<commit_msg>[tensorflow/compiler/xla/tests/select_test.cc] Add calls to `reserve()` before populating vectors<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.
==============================================================================*/

#include <memory>
#include <vector>

#include "tensorflow/compiler/xla/client/global_data.h"
#include "tensorflow/compiler/xla/client/local_client.h"
#include "tensorflow/compiler/xla/client/xla_builder.h"
#include "tensorflow/compiler/xla/tests/client_library_test_base.h"
#include "tensorflow/compiler/xla/tests/literal_test_util.h"
#include "tensorflow/compiler/xla/tests/test_macros.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/types.h"

namespace xla {
namespace {

class SelectTest : public ClientLibraryTestBase {
 public:
  ErrorSpec error_spec_{0.0001};
};

TEST_F(SelectTest, SelectScalarF32True) {
  XlaBuilder builder(TestName());
  auto pred = ConstantR0<bool>(&builder, true);
  auto on_true = ConstantR0<float>(&builder, 123.0f);
  auto on_false = ConstantR0<float>(&builder, 42.0f);
  Select(pred, on_true, on_false);

  ComputeAndCompareR0<float>(&builder, 123.0f, {}, error_spec_);
}

TEST_F(SelectTest, SelectScalarS32True) {
  XlaBuilder builder(TestName());
  auto pred = ConstantR0<bool>(&builder, true);
  auto on_true = ConstantR0<int32>(&builder, -42);
  auto on_false = ConstantR0<int32>(&builder, 42);
  Select(pred, on_true, on_false);

  ComputeAndCompareR0<int32>(&builder, -42, {});
}

TEST_F(SelectTest, SelectScalarF32False) {
  XlaBuilder builder(TestName());
  auto pred = ConstantR0<bool>(&builder, false);
  auto on_true = ConstantR0<float>(&builder, 123.0f);
  auto on_false = ConstantR0<float>(&builder, 42.0f);
  Select(pred, on_true, on_false);

  ComputeAndCompareR0<float>(&builder, 42.0f, {}, error_spec_);
}

XLA_TEST_F(SelectTest, SelectR1S0F32WithConstantR1S0PRED) {
  XlaBuilder builder(TestName());
  auto pred = ConstantR1<bool>(&builder, {});
  auto on_true = ConstantR1<float>(&builder, {});
  auto on_false = ConstantR1<float>(&builder, {});
  Select(pred, on_true, on_false);

  ComputeAndCompareR1<float>(&builder, {}, {}, error_spec_);
}

TEST_F(SelectTest, SelectR1F32WithConstantR1PRED) {
  XlaBuilder builder(TestName());
  auto pred = ConstantR1<bool>(&builder, {false, true, false, true, false});
  auto on_true =
      ConstantR1<float>(&builder, {-2.5f, 25.5f, 2.25f, -10.0f, 6.0f});
  auto on_false =
      ConstantR1<float>(&builder, {10.0f, 5.0f, 1.0f, 10.0f, -6.0f});
  Select(pred, on_true, on_false);

  ComputeAndCompareR1<float>(&builder, {10.0f, 25.5f, 1.0f, -10.0f, -6.0f}, {},
                             error_spec_);
}

XLA_TEST_F(SelectTest, SelectR1S0F32WithCmpR1S0S32s) {
  // Similar to SelectR1S0F32WithConstantR1S0PRED, except that the pred vector
  // is not a constant, but rather the result of comparing two other vectors.
  XlaBuilder builder(TestName());
  auto v1 = ConstantR1<int32>(&builder, {});
  auto v2 = ConstantR1<int32>(&builder, {});
  auto cmp = Eq(v1, v2);
  auto on_true = ConstantR1<float>(&builder, {});
  auto on_false = ConstantR1<float>(&builder, {});
  Select(cmp, on_true, on_false);

  ComputeAndCompareR1<float>(&builder, {}, {}, error_spec_);
}

TEST_F(SelectTest, SelectR1F32WithCmpR1S32s) {
  // Similar to SelectR1F32WithConstantR1PRED, except that the pred vector is
  // not a constant, but rather the result of comparing two other vectors.
  XlaBuilder builder(TestName());
  auto v1 = ConstantR1<int32>(&builder, {1, 2, 3, 4, 5});
  auto v2 = ConstantR1<int32>(&builder, {9, 2, 9, 4, 9});
  auto cmp = Eq(v1, v2);
  auto on_true =
      ConstantR1<float>(&builder, {-2.5f, 25.5f, 2.25f, -10.0f, 6.0f});
  auto on_false =
      ConstantR1<float>(&builder, {10.0f, 5.0f, 1.0f, 10.0f, -6.0f});
  Select(cmp, on_true, on_false);

  ComputeAndCompareR1<float>(&builder, {10.0f, 25.5f, 1.0f, -10.0f, -6.0f}, {},
                             error_spec_);
}

TEST_F(SelectTest, SelectR1F32WithCmpR1F32s) {
  // Similar to SelectR1F32WithCmpR1S32s, except "gt"-comparing two R1F32s.
  XlaBuilder builder(TestName());
  auto v1 = ConstantR1<float>(&builder, {1.0f, 2.0f, 3.0f, 4.0f, 5.0f});
  auto v2 = ConstantR1<float>(&builder, {-1.0f, -2.0f, 13.0f, 14.0f, 4.4f});
  auto cmp = Gt(v1, v2);
  auto on_true =
      ConstantR1<float>(&builder, {-2.5f, 25.5f, 2.25f, -10.0f, 6.0f});
  auto on_false =
      ConstantR1<float>(&builder, {10.0f, 5.0f, 1.0f, 10.0f, -6.0f});
  Select(cmp, on_true, on_false);

  ComputeAndCompareR1<float>(&builder, {-2.5f, 25.5f, 1.0f, 10.0f, 6.0f}, {},
                             error_spec_);
}

TEST_F(SelectTest, SelectR1F32WithCmpR1F32sFromParamsSmall) {
  // Selects among two R1F32s, which come from parameters. v1 and v2 are
  // compared, and selection between them happens based on a gt-comparison mask.
  XlaBuilder builder(TestName());

  XlaOp v1, v2;
  std::unique_ptr<GlobalData> param0_data = CreateR1Parameter<float>(
      {41.0f, 2.0f, 3.0f, 84.0f}, /*parameter_number=*/0, /*name=*/"v1",
      /*builder=*/&builder, /*data_handle=*/&v1);
  std::unique_ptr<GlobalData> param1_data = CreateR1Parameter<float>(
      {21.0f, 22.0f, 23.0f, 24.0f}, /*parameter_number=*/1, /*name=*/"v2",
      /*builder=*/&builder, /*data_handle=*/&v2);

  auto cmp = Gt(v1, v2);
  Select(cmp, v1, v2);
  ComputeAndCompareR1<float>(&builder, {41.0f, 22.0f, 23.0f, 84.0f},
                             {param0_data.get(), param1_data.get()},
                             error_spec_);
}

TEST_F(SelectTest, SelectR1F32WithCmpR1F32sFromParamsLarge) {
  // Similar to SelectR1F32WithCmpR1F32sFromParamsSmall, except that the
  // data size passed in and out is large.
  XlaBuilder builder(TestName());

  // Number of floats in the data passed into and out of the computation.
  constexpr int datalen = 15 * 1000;

  // The inputs are initialized with a special pattern where in the first third
  // of the data v1[i] > v2[i] and elsewhere it's vice versa.
  std::vector<float> v1vec;
  std::vector<float> v2vec;
  std::vector<float> expected_vec;
  v1vec.reserve(datalen);
  v2vec.reserve(datalen);
  expected_vec.reserve(datalen);
  for (int i = 0; i < datalen; ++i) {
    float smaller = i;
    float larger = i * 2;
    if (i < datalen / 3) {
      v1vec.push_back(larger);
      v2vec.push_back(smaller);
    } else {
      v1vec.push_back(smaller);
      v2vec.push_back(larger);
    }
    expected_vec.push_back(larger);
  }

  XlaOp v1, v2;
  std::unique_ptr<GlobalData> param0_data =
      CreateR1Parameter<float>(v1vec, /*parameter_number=*/0, /*name=*/"v1",
                               /*builder=*/&builder, /*data_handle=*/&v1);
  std::unique_ptr<GlobalData> param1_data =
      CreateR1Parameter<float>(v2vec, /*parameter_number=*/1, /*name=*/"v2",
                               /*builder=*/&builder, /*data_handle=*/&v2);

  auto cmp = Gt(v1, v2);
  Select(cmp, v1, v2);
  ComputeAndCompareR1<float>(&builder, expected_vec,
{'fraction_non_alphanumeric': 0.15887472171625178, 'fraction_numerical': 0.05585913782635094, 'mean_word_length': 4.047327204630576, 'pattern_counts': {'":': 0, '<': 107, '<?xml version=': 0, '>': 105, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2225751', 'n_tokens_mistral': 5830, 'n_tokens_neox': 5491, 'n_words': 1197}