/* 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 <stdint.h>

#include <cstdlib>
#include <cstring>
#include <initializer_list>
#include <string>
#include <type_traits>
#include <vector>

#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "Eigen/Core"  // from @eigen_archive
#include "tensorflow/lite/c/c_api_types.h"
#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/kernels/internal/portable_tensor_utils.h"
#include "tensorflow/lite/kernels/test_util.h"
#include "tensorflow/lite/schema/schema_generated.h"
#include "tensorflow/lite/types/half.h"

namespace tflite {
namespace {

using ::testing::ElementsAreArray;

template <typename InputType, typename PositionsType>
class GatherOpModel : public SingleOpModel {
 public:
  GatherOpModel(const TensorData& input, const TensorData& positions,
                bool constant_tensor, const std::vector<InputType>& input_data,
                const std::vector<PositionsType>& positions_data, int axis = 0,
                int batch_dims = 0) {
    if (constant_tensor) {
      input_ = AddConstInput(input, input_data);
      positions_ = AddConstInput(positions, positions_data);
    } else {
      input_ = AddInput(input);
      positions_ = AddInput(positions);
    }
    output_ = AddOutput(input.type);
    SetBuiltinOp(BuiltinOperator_GATHER, BuiltinOptions_GatherOptions,
                 CreateGatherOptions(builder_, axis, batch_dims).Union());
    BuildInterpreter({GetShape(input_), GetShape(positions_)});
    if (!constant_tensor) {
      if (input.type == TensorType_INT4) {
        SetInputInt4(input_, input_data,
                     std::is_same<std::string, InputType>());
      } else {
        SetInput(input_, input_data, std::is_same<std::string, InputType>());
      }
      SetPositions(positions_data);
    }
  }

  template <typename T>
  void SetInput(int input, const std::vector<T> data, std::false_type) {
    PopulateTensor<T>(input, data);
  }

  // Overload for string inputs.
  template <typename T>
  void SetInput(int input, const std::vector<T> data, std::true_type) {
    PopulateStringTensor(input_, data);
  }

  template <typename T>
  void SetInputInt4(int input, const std::vector<T> data, std::false_type) {
    auto non_const = *const_cast<std::vector<T>*>(&data);
    std::vector<int8_t> data_int8(non_const.size());
    std::copy(non_const.begin(), non_const.end(), data_int8.begin());
    PopulateTensor4bit(input, 0, data_int8.data(),
                       data_int8.data() + data_int8.size());
  }

  template <typename T>
  void SetInputInt4(int input, const std::vector<T> data, std::true_type) {
    // Unsupported
  }

  void SetPositions(const std::vector<PositionsType>& data) {
    PopulateTensor<PositionsType>(positions_, data);
  }

  std::vector<InputType> GetOutput() {
    return ExtractVector<InputType>(output_);
  }

  std::vector<std::string> GetStringOutput() {
    return ExtractVector<std::string>(output_);
  }

  std::vector<int8_t> GetInt4Output() {
    const auto* tensor = interpreter_->tensor(output_);
    const std::vector<int8_t> data_int8 = std::vector<int8_t>(
        tensor->data.raw, tensor->data.raw + GetTensorSize(output_));
    int num_elements = 1;
    auto shape = GetTensorShape(output_);
    for (int i = 0; i < shape.size(); i++) {
      num_elements *= shape[i];
    }
    std::vector<int8_t> inflated_output(num_elements);
    tensor_utils::UnpackPackedIntToInt8(data_int8.data(), num_elements,
                                        /*bit_width=*/4,
                                        inflated_output.data());
    return inflated_output;
  }

  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }

  void SetRawInput(const char* data, size_t bytes) {
    auto tensor = interpreter_->tensor(input_);
    char* tensor_buffer = reinterpret_cast<char*>(malloc(bytes));
    memcpy(tensor_buffer, data, bytes);
    TfLiteTensorReset(tensor->type, tensor->name,
                      TfLiteIntArrayCopy(tensor->dims), tensor->params,
                      tensor_buffer, bytes, kTfLiteDynamic, tensor->allocation,
                      tensor->is_variable, tensor);
  }

 protected:
  int input_;
  int positions_;
  int output_;
};

struct GatherOpTest : public testing::TestWithParam<bool> {};

INSTANTIATE_TEST_SUITE_P(ConstantTensor, GatherOpTest, testing::Bool());

TEST_P(GatherOpTest, Shuffle) {
  bool constant_tensor = GetParam();
  GatherOpModel<float, int32_t> m({TensorType_FLOAT32, {2, 2}},
                                  {TensorType_INT32, {2}}, constant_tensor,
                                  {-2.0, 0.2, 0.7, 0.8}, {1, 0});
  ASSERT_EQ(m.Invoke(), kTfLiteOk);
  EXPECT_THAT(m.GetOutput(),
              ElementsAreArray(ArrayFloatNear({0.7, 0.8, -2, 0.2})));
}

TEST_P(GatherOpTest, Test0DIndex) {
  bool constant_tensor = GetParam();
  GatherOpModel<float, int32_t> m({TensorType_FLOAT32, {2, 2}},
                                  {TensorType_INT32, {}}, constant_tensor,
                                  {-2.0, 0.2, 0.7, 0.8}, {1});
  ASSERT_EQ(m.Invoke(), kTfLiteOk);
  EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({0.7, 0.8})));
  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2}));
}

TEST_P(GatherOpTest, Test0DIndexWith0DResult) {
  bool constant_tensor = GetParam();
  // 0D tensor is special case in current TFLite. Test it once to make sure
  // existing workarounds are fine with it.
  GatherOpModel<float, int32_t> m({TensorType_FLOAT32, {3}},
                                  {TensorType_INT32, {}}, constant_tensor,
                                  {1.0, 2.0, 3.0}, {1});
  ASSERT_EQ(m.Invoke(), kTfLiteOk);
  EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({2.0})));
  EXPECT_TRUE(m.GetOutputShape().empty());
}

TEST_P(GatherOpTest, Test1DInput1DIndex) {
  bool constant_tensor = GetParam();
  GatherOpModel<float, int32_t> m({TensorType_FLOAT32, {3}},
                                  {TensorType_INT32, {1}}, constant_tensor,
                                  {1.0, 3.0, 5.0}, {1});
  ASSERT_EQ(m.Invoke(), kTfLiteOk);
  EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({3.0})));
  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1}));
}

TEST_P(GatherOpTest, Test2DIndexWith2DResult) {
  bool constant_tensor = GetParam();
  GatherOpModel<float, int32_t> m({TensorType_FLOAT32, {3}},
                                  {TensorType_INT32, {1, 2}}, constant_tensor,
                                  {1.0, 2.0, 3.0}, {1, 0});
  ASSERT_EQ(m.Invoke(), kTfLiteOk);
  EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({2.0, 1.0})));
  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2}));
}

TEST_P(GatherOpTest, Duplicate) {
  bool constant_tensor = GetParam();
  GatherOpModel<float, int32_t> m({TensorType_FLOAT32, {1, 2, 2}},
                                  {TensorType_INT32, {2}}, constant_tensor,
                                  {-2.0, 0.2, 0.7, 0.8}, {0, 0});
  ASSERT_EQ(m.Invoke(), kTfLiteOk);
  EXPECT_THAT(
      m.GetOutput(),
      ElementsAreArray(ArrayFloatNear({-2, 0.2, 0.7, 0.8, -2, 0.2, 0.7, 0.8})));
}

TEST_P(GatherOpTest, Slice) {
  bool constant_tensor = GetParam();
  GatherOpModel<float, int32_t> m({TensorType_FLOAT32, {4, 1}},
                                  {TensorType_INT32, {2}}, constant_tensor,
                                  {-2.0, 0.2, 0.7, 0.8}, {1, 3});
  ASSERT_EQ(m.Invoke(), kTfLiteOk);
  EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({0.2, 0.8})));
}

TEST_P(GatherOpTest, Axis1) {
  bool constant_tensor = GetParam();
  const int axis = 1;
  GatherOpModel<float, int32_t> m({TensorType_FLOAT32, {1, 2, 3}},
                                  {TensorType_INT32, {2}}, constant_tensor,
                                  {1, 2, 3, 4, 5, 6}, {1, 0}, axis);
  ASSERT_EQ(m.Invoke(), kTfLiteOk);
  EXPECT_THAT(m.GetOutput(),
              ElementsAreArray(ArrayFloatNear({4, 5, 6, 1, 2, 3})));
  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2, 3}));
}

TEST_P(GatherOpTest, Axis10DIndex) {
  bool constant_tensor = GetParam();
  const int axis = 1;
  GatherOpModel<float, int32_t> m({TensorType_FLOAT32, {1, 3, 2}},
                                  {TensorType_INT32, {}}, constant_tensor,
                                  {1, 2, 3, 4, 5, 6}, {1}, axis);
  ASSERT_EQ(m.Invoke(), kTfLiteOk);
  EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({3, 4})));
  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2}));
}

TEST_P(GatherOpTest, Axis1Slice) {
  bool constant_tensor = GetParam();
  const int axis = 1;
  GatherOpModel<float, int32_t> m({TensorType_FLOAT32, {1, 4, 2}},
                                  {TensorType_INT32, {2}}, constant_tensor,
                                  {1, 2, 3, 4, 5, 6, 7, 8}, {3, 1}, axis);
  ASSERT_EQ(m.Invoke(), kTfLiteOk);
  EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({7, 8, 3, 4})));
  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2, 2}));
}

TEST_P(GatherOpTest, LastAxis) {
  const int axis = -1;
  bool constant_tensor = GetParam();
  GatherOpModel<float, int32_t> m({TensorType_FLOAT32, {1, 2, 3}},
                                  {TensorType_INT32, {2}}, constant_tensor,
                                  {1, 2, 3, 4, 5, 6}, {2, 0}, axis);
  ASSERT_EQ(m.Invoke(), kTfLiteOk);
  EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({3, 1, 6, 4})));
  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2, 2}));
}

TEST_P(GatherOpTest, LastAxis0DIndex) {
  bool constant_tensor = GetParam();
  const int axis = -1;
  GatherOpModel<float, int32_t> m({TensorType_FLOAT32, {1, 2, 3}},
                                  {TensorType_INT32, {}}, constant_tensor,
                                  {1, 2, 3, 4, 5, 6}, {2}, axis);
  ASSERT_EQ(m.Invoke(), kTfLiteOk);
  EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({3, 6})));
  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2}));
}

using TestTypes = testing::Types<int8_t, uint8_t, int16_t, int32_t, int64_t,
                                 float, half, Eigen::bfloat16>;

template <typename T>
struct TypedGatherOpTest : public testing::Test {};

TYPED_TEST_CASE(TypedGatherOpTest, TestTypes);

TYPED_TEST(TypedGatherOpTest, Int32Indices) {
  for (bool constant_tensor : {true, false}) {
    TensorType tensor_type = GetTensorType<TypeParam>();
    GatherOpModel<TypeParam, int32_t> m(
        {tensor_type, {2, 2}}, {TensorType_INT32, {2}}, constant_tensor,
        {TypeParam(13), TypeParam(120), TypeParam(14), TypeParam(15)}, {1, 0});
    ASSERT_EQ(m.Invoke(), kTfLiteOk);

    EXPECT_THAT(m.GetOutput(),
                ElementsAreArray({TypeParam(14), TypeParam(15), TypeParam(13),
                                  TypeParam(120)}));
  }
}

TYPED_TEST(TypedGatherOpTest, Int64Indices) {
  for (bool constant_tensor : {true, false}) {
    TensorType tensor_type = GetTensorType<TypeParam>();
    GatherOpModel<TypeParam, int64_t> m(
        {tensor_type, {2, 2}}, {TensorType_INT64, {2}}, constant_tensor,
        {TypeParam(13), TypeParam(120), TypeParam(14), TypeParam(15)}, {1, 0});
    ASSERT_EQ(m.Invoke(), kTfLiteOk);

    EXPECT_THAT(m.GetOutput(),
                ElementsAreArray({TypeParam(14), TypeParam(15), TypeParam(13),
                                  TypeParam(120)}));
  }
}

TEST(GatherOpTest, SimpleString) {
  GatherOpModel<std::string, int32_t> m(
      {TensorType_STRING, {3}}, {TensorType_INT32, {2}},
      /*constant_tensor=*/false, {"A", "B", "C"}, {0, 2});
  ASSERT_EQ(m.Invoke(), kTfLiteOk);
  ASSERT_THAT(m.GetOutputShape(), ElementsAreArray({2}));
  EXPECT_THAT(m.GetStringOutput(), ElementsAreArray({"A", "C"}));
}

TEST(GatherOpTest, StringIndexTruncation) {
  GatherOpModel<std::string, int16_t> m({TensorType_STRING, {1}},
                                        {TensorType_INT16, {1}},
                                        /*constant_tensor=*/false, {"A"}, {0});

  // Access the implementation details to manually corrupt the string tensor's
  // buffer. We want to simulate:
  // - num_strings = -65535 (which is 0xFFFF0001, truncates to 1 in int16_t)
  // - indexes = {0}
  // - pos = 0 < 1 check would pass in 16-bit, but should fail with our
  // validation.

  int32_t malformed_data[3];
  malformed_data[0] = -65535;  // N
  malformed_data[1] = 12;      // offset
  malformed_data[2] = 12;      // total length

  m.SetRawInput(reinterpret_cast<const char*>(malformed_data),
                sizeof(malformed_data));

  // Invoke should fail (not kTfLiteOk)
  EXPECT_NE(m.Invoke(), kTfLiteOk);
}

TEST_P(GatherOpTest, 2DIndexString) {
  GatherOpModel<std::string, int32_t> m(
      {TensorType_STRING, {3}}, {TensorType_INT32, {2, 3}},
      /*constant_tensor=*/false, {"A", "B", "C"}, {0, 2, 1, 1, 0, 2});
  ASSERT_EQ(m.Invoke(), kTfLiteOk);
  ASSERT_THAT(m.GetOutputShape(), ElementsAreArray({2, 3}));
  EXPECT_THAT(m.GetStringOutput(),
              ElementsAreArray({"A", "C", "B", "B", "A", "C"}));
}

TYPED_TEST(TypedGatherOpTest, BatchDims2) {
  for (bool constant_tensor : {true, false}) {
    TensorType tensor_type = GetTensorType<TypeParam>();
    GatherOpModel<TypeParam, int32_t> m(
        {tensor_type, {2, 2, 3, 5}}, {TensorType_INT32, {2, 2, 2}},
        constant_tensor,
        {TypeParam(0),  TypeParam(1),  TypeParam(2),  TypeParam(3),
         TypeParam(4),  TypeParam(5),  TypeParam(6),  TypeParam(7),
         TypeParam(8),  TypeParam(9),  TypeParam(10), TypeParam(11),
         TypeParam(12), TypeParam(13), TypeParam(14), TypeParam(15),
         TypeParam(16), TypeParam(17), TypeParam(18), TypeParam(19),
         TypeParam(20), TypeParam(21), TypeParam(22), TypeParam(23),
         TypeParam(24), TypeParam(25), TypeParam(26), TypeParam(27),
         TypeParam(28), TypeParam(29), TypeParam(30), TypeParam(31),
         TypeParam(32), TypeParam(33), TypeParam(34), TypeParam(35),
         TypeParam(36), TypeParam(37), TypeParam(38), TypeParam(39),
         TypeParam(40), TypeParam(41), TypeParam(42), TypeParam(43),
         TypeParam(44), TypeParam(45), TypeParam(46), TypeParam(47),
         TypeParam(48), TypeParam(49), TypeParam(50), TypeParam(51),
         TypeParam(52), TypeParam(53), TypeParam(54), TypeParam(55),
         TypeParam(56), TypeParam(57), TypeParam(58), TypeParam(59)},
        {1, 0, 0, 1, 1, 0, 0, 1},
        /*axis=*/2,
        /*batch_dims=*/2);
    ASSERT_EQ(m.Invoke(), kTfLiteOk);

    ASSERT_THAT(m.GetOutputShape(), ElementsAreArray({2, 2, 2, 5}));
    EXPECT_THAT(
        m.GetOutput(),
        ElementsAreArray(
            {TypeParam(5),  TypeParam(6),  TypeParam(7),  TypeParam(8),
             TypeParam(9),  TypeParam(0),  TypeParam(1),  TypeParam(2),
             TypeParam(3),  TypeParam(4),  TypeParam(15), TypeParam(16),
             TypeParam(17), TypeParam(18), TypeParam(19), TypeParam(20),
             TypeParam(21), TypeParam(22), TypeParam(23), TypeParam(24),
             TypeParam(35), TypeParam(36), TypeParam(37), TypeParam(38),
             TypeParam(39), TypeParam(30), TypeParam(31), TypeParam(32),
             TypeParam(33), TypeParam(34), TypeParam(45), TypeParam(46),
             TypeParam(47), TypeParam(48), TypeParam(49), TypeParam(50),
             TypeParam(51), TypeParam(52), TypeParam(53), TypeParam(54)}));
  }
}

TEST_P(GatherOpTest, BatchDims1) {
  bool constant_tensor = GetParam();
  GatherOpModel<int8_t, int32_t> m(
      {TensorType_INT8, {2, 2, 3, 5}}, {TensorType_INT32, {2, 2, 2}},
      constant_tensor,
      {0,  1,  2,  3,  4,  5,  6,  7,  8,  9,  10, 11, 12, 13, 14,
       15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
       30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44,
       45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59},
      {1, 0, 0, 1, 1, 0, 0, 1},
      /*axis=*/2, /*batch_dims=*/1);
  ASSERT_EQ(m.Invoke(), kTfLiteOk);

  ASSERT_THAT(m.GetOutputShape(), ElementsAreArray({2, 2, 2, 2, 5}));
  EXPECT_THAT(
      m.GetOutput(),
      ElementsAreArray({5,  6,  7,  8,  9,  0,  1,  2,  3,  4,  0,  1,  2,  3,
                        4,  5,  6,  7,  8,  9,  20, 21, 22, 23, 24, 15, 16, 17,
                        18, 19, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 35, 36,
                        37, 38, 39, 30, 31, 32, 33, 34, 30, 31, 32, 33, 34, 35,
                        36, 37, 38, 39, 50, 51, 52, 53, 54, 45, 46, 47, 48, 49,
                        45, 46, 47, 48, 49, 50, 51, 52, 53, 54}));
}

TEST_P(GatherOpTest, NegativeBatchDims) {
  bool constant_tensor = GetParam();
  GatherOpModel<int8_t, int32_t> m(
      {TensorType_INT8, {2, 2, 3, 5}}, {TensorType_INT32, {2, 2, 2}},
      constant_tensor,
      {0,  1,  2,  3,  4,  5,  6,  7,  8,  9,  10, 11, 12, 13, 14,
       15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
       30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44,
       45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59},
      {1, 0, 0, 1, 1, 0, 0, 1},
      /*axis=*/2, /*batch_dims=*/-2);
  ASSERT_EQ(m.Invoke(), kTfLiteOk);

  ASSERT_THAT(m.GetOutputShape(), ElementsAreArray({2, 2, 2, 2, 5}));
  EXPECT_THAT(
      m.GetOutput(),
      ElementsAreArray({5,  6,  7,  8,  9,  0,  1,  2,  3,  4,  0,  1,  2,  3,
                        4,  5,  6,  7,  8,  9,  20, 21, 22, 23, 24, 15, 16, 17,
                        18, 19, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 35, 36,
                        37, 38, 39, 30, 31, 32, 33, 34, 30, 31, 32, 33, 34, 35,
                        36, 37, 38, 39, 50, 51, 52, 53, 54, 45, 46, 47, 48, 49,
                        45, 46, 47, 48, 49, 50, 51, 52, 53, 54}));
}

TEST_P(GatherOpTest, BatchDimsEqualIndexDims) {
  bool constant_tensor = GetParam();
  GatherOpModel<int8_t, int32_t> m(
      {TensorType_INT8, {2, 2, 2, 5}}, {TensorType_INT32, {2, 2, 2}},
      constant_tensor, {0,  1,  2,  3,  4,  5,  6,  7,  8,  9,  10, 11, 12, 13,
                        14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27,
                        28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39},
      {1, 0, 0, 1, 1, 0, 0, 1},
      /*axis=*/3, /*batch_dims=*/3);
  ASSERT_EQ(m.Invoke(), kTfLiteOk);

  ASSERT_THAT(m.GetOutputShape(), ElementsAreArray({2, 2, 2}));
  EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 5, 10, 16, 21, 25, 30, 36}));
}

TEST_P(GatherOpTest, ErrorOnOutOfBoundsTooLarge) {
  bool constant_tensor = GetParam();
  if (constant_tensor) {
#if GTEST_HAS_DEATH_TEST
    EXPECT_DEATH(
        (GatherOpModel<float, int32_t>({TensorType_FLOAT32, {2, 2}},
                                       {TensorType_INT32, {2}}, constant_tensor,
                                       {
                                           -2.f, 0.2f,  //
                                           0.7f, 0.8f   //
                                       },
                                       {3, 1})),
        "Cannot allocate tensors");
#endif
  } else {
    GatherOpModel<float, int32_t> m({TensorType_FLOAT32, {2, 2}},
                                    {TensorType_INT32, {2}}, constant_tensor,
                                    {
                                        -2.f, 0.2f,  //
                                        0.7f, 0.8f   //
                                    },
                                    {3, 1});
    EXPECT_EQ(m.Invoke(), kTfLiteError);
  }
}

TEST_P(GatherOpTest, ErrorOnOutOfBoundsNegative) {
  bool constant_tensor = GetParam();
  if (constant_tensor) {
#if GTEST_HAS_DEATH_TEST
    EXPECT_DEATH(
        (GatherOpModel<float, int32_t>({TensorType_FLOAT32, {2, 2}},
                                       {TensorType_INT32, {2}}, constant_tensor,
                                       {
                                           -2.f, 0.2f,  //
                                           0.7f, 0.8f   //
                                       },
                                       {-1, 0})),
        "Cannot allocate tensors");
#endif
  } else {
    GatherOpModel<float, int32_t> m({TensorType_FLOAT32, {2, 2}},
                                    {TensorType_INT32, {2}}, constant_tensor,
                                    {
                                        -2.f, 0.2f,  //
                                        0.7f, 0.8f   //
                                    },
                                    {-1, 0});
    ASSERT_EQ(m.Invoke(), kTfLiteError);
    m.SetPositions({-1, 0});
    EXPECT_EQ(m.Invoke(), kTfLiteError);
  }
}

TEST(GatherOpTest, BatchDims1Int4) {
  GatherOpModel<int8_t, int32_t> m(
      {TensorType_INT4, {2, 2, 3, 4}}, {TensorType_INT32, {2, 2, 2}}, false,
      {1,  2,  3,  4,  -1, -2, -3, -4, 0,  0,  0,  0,  1,  2,  3,  4,
       -1, -2, -3, -4, 0,  0,  0,  0,  4,  5,  6,  7,  -5, -6, -7, -8,
       0,  0,  0,  0,  4,  5,  6,  7,  -5, -6, -7, -8, 0,  0,  0,  0},
      {1, 0, 0, 1, 1, 0, 0, 1},
      /*axis=*/2, /*batch_dims=*/1);
  ASSERT_EQ(m.Invoke(), kTfLiteOk);

  ASSERT_THAT(m.GetOutputShape(), ElementsAreArray({2, 2, 2, 2, 4}));

  EXPECT_THAT(m.GetInt4Output(),
              ElementsAreArray(
                  {-1, -2, -3, -4, 1, 2, 3, 4, 1, 2, 3, 4, -1, -2, -3, -4,
                   -1, -2, -3, -4, 1, 2, 3, 4, 1, 2, 3, 4, -1, -2, -3, -4,
                   -5, -6, -7, -8, 4, 5, 6, 7, 4, 5, 6, 7, -5, -6, -7, -8,
                   -5, -6, -7, -8, 4, 5, 6, 7, 4, 5, 6, 7, -5, -6, -7, -8}));
}

}  // namespace
}  // namespace tflite