// Copyright (c) ONNX Project Contributors
//
// SPDX-License-Identifier: Apache-2.0

#include "onnx_core/backend_test/test_case.h"
#include "onnx_core/compute/raw_buffer_allocator.h"
#include "onnx_core/runtime/kernels/cast_helper.h"
#include "onnx_core/runtime/kernels/kernel_context.h"
#include "onnx_core/runtime/kernels/parallel_for.h"
#include "onnx_core/runtime/runtime_context.h"
#include "onnx_core/runtime/tuning/kernel_tuning.h"
#include "onnx_extensions/kernels/kernels/logical/include_logical_kernels.h"

#include <gtest/gtest.h>

#include <limits>
#include <stdexcept>
#include <vector>

using namespace ONNX_LIGHT_NAMESPACE;
using core::backend_test::DefaultOpset;
using core::runtime::RuntimeContext;
using core::runtime::Tensor;
using onnx_kernels::SimpleRawBufferAllocator;
using onnx_kernels::kernel::And;
using onnx_kernels::kernel::BitwiseAnd;
using onnx_kernels::kernel::BitwiseNot;
using onnx_kernels::kernel::BitwiseOr;
using onnx_kernels::kernel::BitwiseXor;
using onnx_kernels::kernel::Equal;
using onnx_kernels::kernel::Greater;
using onnx_kernels::kernel::GreaterOrEqual;
using onnx_kernels::kernel::IsInf;
using onnx_kernels::kernel::IsNaN;
using onnx_kernels::kernel::KernelContext;
using onnx_kernels::kernel::Less;
using onnx_kernels::kernel::LessOrEqual;
using onnx_kernels::kernel::Not;
using onnx_kernels::kernel::Or;
using onnx_kernels::kernel::Where;
using onnx_kernels::kernel::Xor;

namespace Test {

TEST(KernelClass, AndClassMatchesReference) {
  const KernelContext ctx{DefaultOpset(7)};
  And and_kernel{ctx};
  Tensor x("", core::runtime::DataType::BOOL, {2, 2}, {1, 0, 1, 0});
  Tensor y("", core::runtime::DataType::BOOL, {2, 2}, {1, 1, 0, 0});
  Tensor z = and_kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  EXPECT_EQ(z.data_type, static_cast<int32_t>(core::runtime::DataType::BOOL));
  EXPECT_EQ(z.data[0], 1);
  EXPECT_EQ(z.data[1], 0);
  EXPECT_EQ(z.data[2], 0);
  EXPECT_EQ(z.data[3], 0);
}

TEST(KernelClass, AndClassBroadcastsScalar) {
  const KernelContext ctx{DefaultOpset(7)};
  And and_kernel{ctx};
  Tensor x("", core::runtime::DataType::BOOL, {2, 2}, {1, 0, 1, 0});
  Tensor y("", core::runtime::DataType::BOOL, {}, {1});
  Tensor z = and_kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  EXPECT_EQ(z.data[0], 1);
  EXPECT_EQ(z.data[1], 0);
  EXPECT_EQ(z.data[2], 1);
  EXPECT_EQ(z.data[3], 0);
}

TEST(KernelClass, AndUsesTypedParallelTuningForBroadcasting) {
  const KernelContext ctx{DefaultOpset(7)};
  And and_kernel{ctx};
  const auto bool_key = and_kernel.TuningKey(static_cast<int32_t>(core::runtime::DataType::BOOL));

  EXPECT_EQ(bool_key.kernel, "And");
  EXPECT_EQ(and_kernel.TuningKey(static_cast<int32_t>(core::runtime::DataType::FLOAT)).device,
            core::symbolic::Device::kUndefined);
  const auto schema = core::runtime::GetKernelTuningRegistry().FindSchema(bool_key);
  ASSERT_NE(schema, nullptr);
  EXPECT_EQ(schema->portable_defaults().Get<int64_t>("parallel.minimum_elements"),
            core::runtime::kParallelForGrainSize);

  and_kernel.Configure({bool_key, {{"parallel.minimum_elements", int64_t{1}}}});
  EXPECT_EQ(and_kernel.tuning().parallel_minimum_elements, 1);
  Tensor x("", core::runtime::DataType::BOOL, {2, 1}, {1, 0});
  Tensor y("", core::runtime::DataType::BOOL, {1, 3}, {1, 0, 1});
  Tensor z = and_kernel(x, y);
  EXPECT_EQ(z.shape, (std::vector<int64_t>{2, 3}));
  EXPECT_EQ(z.data, (std::vector<uint8_t>{1, 0, 1, 0, 0, 0}));
}

TEST(KernelClass, OrClassMatchesReference) {
  const KernelContext ctx{DefaultOpset(7)};
  Or or_kernel{ctx};
  Tensor x("", core::runtime::DataType::BOOL, {2, 2}, {1, 0, 1, 0});
  Tensor y("", core::runtime::DataType::BOOL, {2, 2}, {1, 1, 0, 0});
  Tensor z = or_kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  EXPECT_EQ(z.data_type, static_cast<int32_t>(core::runtime::DataType::BOOL));
  EXPECT_EQ(z.data[0], 1);
  EXPECT_EQ(z.data[1], 1);
  EXPECT_EQ(z.data[2], 1);
  EXPECT_EQ(z.data[3], 0);
}

TEST(KernelClass, XorClassMatchesReference) {
  const KernelContext ctx{DefaultOpset(7)};
  Xor xor_kernel{ctx};
  Tensor x("", core::runtime::DataType::BOOL, {2, 2}, {1, 0, 1, 0});
  Tensor y("", core::runtime::DataType::BOOL, {2, 2}, {1, 1, 0, 0});
  Tensor z = xor_kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  EXPECT_EQ(z.data_type, static_cast<int32_t>(core::runtime::DataType::BOOL));
  EXPECT_EQ(z.data[0], 0);
  EXPECT_EQ(z.data[1], 1);
  EXPECT_EQ(z.data[2], 1);
  EXPECT_EQ(z.data[3], 0);
}

TEST(KernelClass, LogicalRejectsNonBoolTensors) {
  const KernelContext ctx{DefaultOpset(7)};
  And and_kernel{ctx};
  Tensor x = Tensor::FromFloat("", {2}, {1.0f, 0.0f});
  Tensor y("", core::runtime::DataType::BOOL, {2}, {1, 0});
  EXPECT_THROW((void)and_kernel(x, y), std::invalid_argument);
}

TEST(KernelClass, AndInPlaceWritesToPreallocatedOutput) {
  const KernelContext ctx{DefaultOpset(7)};
  And and_kernel{ctx};
  Tensor x("", core::runtime::DataType::BOOL, {2, 2}, {1, 0, 1, 0});
  Tensor y("", core::runtime::DataType::BOOL, {2, 2}, {1, 1, 0, 0});
  Tensor z("", core::runtime::DataType::BOOL, {2, 2}, std::vector<uint8_t>(4, 9));
  and_kernel(x, y, z);
  EXPECT_EQ(z.data[0], 1);
  EXPECT_EQ(z.data[1], 0);
  EXPECT_EQ(z.data[2], 0);
  EXPECT_EQ(z.data[3], 0);
}

TEST(KernelClass, OrInPlaceWritesToPreallocatedOutput) {
  const KernelContext ctx{DefaultOpset(7)};
  Or or_kernel{ctx};
  Tensor x("", core::runtime::DataType::BOOL, {2, 2}, {1, 0, 1, 0});
  Tensor y("", core::runtime::DataType::BOOL, {2, 2}, {1, 1, 0, 0});
  Tensor z("", core::runtime::DataType::BOOL, {2, 2}, std::vector<uint8_t>(4));
  or_kernel(x, y, z);
  EXPECT_EQ(z.data[0], 1);
  EXPECT_EQ(z.data[1], 1);
  EXPECT_EQ(z.data[2], 1);
  EXPECT_EQ(z.data[3], 0);
}

TEST(KernelClass, XorInPlaceWritesToPreallocatedOutput) {
  const KernelContext ctx{DefaultOpset(7)};
  Xor xor_kernel{ctx};
  Tensor x("", core::runtime::DataType::BOOL, {2, 2}, {1, 0, 1, 0});
  Tensor y("", core::runtime::DataType::BOOL, {2, 2}, {1, 1, 0, 0});
  Tensor z("", core::runtime::DataType::BOOL, {2, 2}, std::vector<uint8_t>(4));
  xor_kernel(x, y, z);
  EXPECT_EQ(z.data[0], 0);
  EXPECT_EQ(z.data[1], 1);
  EXPECT_EQ(z.data[2], 1);
  EXPECT_EQ(z.data[3], 0);
}

TEST(KernelClass, GreaterClassMatchesReference) {
  const KernelContext ctx{DefaultOpset(13)};
  Greater greater_kernel{ctx};
  Tensor x = Tensor::FromFloat("", {2, 2}, {1.0f, 2.0f, 3.0f, 4.0f});
  Tensor y = Tensor::FromFloat("", {2, 2}, {2.0f, 2.0f, 2.0f, 2.0f});
  Tensor z = greater_kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  EXPECT_EQ(z.data_type, static_cast<int32_t>(core::runtime::DataType::BOOL));
  EXPECT_EQ(z.data[0], 0); // 1 > 2 -> false
  EXPECT_EQ(z.data[1], 0); // 2 > 2 -> false
  EXPECT_EQ(z.data[2], 1); // 3 > 2 -> true
  EXPECT_EQ(z.data[3], 1); // 4 > 2 -> true
}

TEST(KernelClass, GreaterClassBroadcastsScalar) {
  const KernelContext ctx{DefaultOpset(13)};
  Greater greater_kernel{ctx};
  Tensor x = Tensor::FromFloat("", {2, 2}, {1.0f, 2.0f, 3.0f, 4.0f});
  Tensor y = Tensor::FromFloat("", {}, {2.5f});
  Tensor z = greater_kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  EXPECT_EQ(z.data[0], 0);
  EXPECT_EQ(z.data[1], 0);
  EXPECT_EQ(z.data[2], 1);
  EXPECT_EQ(z.data[3], 1);
}

TEST(KernelClass, GreaterInPlaceWritesToPreallocatedOutput) {
  const KernelContext ctx{DefaultOpset(13)};
  Greater greater_kernel{ctx};
  Tensor x = Tensor::FromFloat("", {2, 2}, {1.0f, 2.0f, 3.0f, 4.0f});
  Tensor y = Tensor::FromFloat("", {2, 2}, {2.0f, 2.0f, 2.0f, 2.0f});
  Tensor z("", core::runtime::DataType::BOOL, {2, 2}, std::vector<uint8_t>(4, 9));
  greater_kernel(x, y, z);
  EXPECT_EQ(z.data[0], 0);
  EXPECT_EQ(z.data[1], 0);
  EXPECT_EQ(z.data[2], 1);
  EXPECT_EQ(z.data[3], 1);
}

TEST(KernelClass, GreaterRejectsUnsupportedDtype) {
  // BOOL inputs are not in the supported dtype set (FLOAT/INT8/INT16/UINT8/
  // UINT16/UINT32/UINT64) so the kernel must reject them.
  const KernelContext ctx{DefaultOpset(13)};
  Greater greater_kernel{ctx};
  Tensor x("", core::runtime::DataType::BOOL, {2}, {1, 0});
  Tensor y("", core::runtime::DataType::BOOL, {2}, {0, 1});
  EXPECT_THROW((void)greater_kernel(x, y), std::invalid_argument);
}

TEST(KernelClass, GreaterClassMatchesReferenceInt8) {
  const KernelContext ctx{DefaultOpset(13)};
  Greater greater_kernel{ctx};
  Tensor x = Tensor::FromInt8("", {4}, {-2, 0, 3, 7});
  Tensor y = Tensor::FromInt8("", {4}, {-1, 0, 1, 9});
  Tensor z = greater_kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  EXPECT_EQ(z.data_type, static_cast<int32_t>(core::runtime::DataType::BOOL));
  EXPECT_EQ(z.data[0], 0); // -2 > -1 -> false
  EXPECT_EQ(z.data[1], 0); //  0 >  0 -> false
  EXPECT_EQ(z.data[2], 1); //  3 >  1 -> true
  EXPECT_EQ(z.data[3], 0); //  7 >  9 -> false
}

TEST(KernelClass, GreaterClassMatchesReferenceUint32) {
  const KernelContext ctx{DefaultOpset(13)};
  Greater greater_kernel{ctx};
  Tensor x = Tensor::FromUint32("", {4}, {1u, 5u, 0u, 7u});
  Tensor y = Tensor::FromUint32("", {4}, {2u, 5u, 0u, 6u});
  Tensor z = greater_kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  EXPECT_EQ(z.data[0], 0);
  EXPECT_EQ(z.data[1], 0);
  EXPECT_EQ(z.data[2], 0);
  EXPECT_EQ(z.data[3], 1);
}

TEST(KernelClass, GreaterClassMatchesReferenceInt64) {
  const KernelContext ctx{DefaultOpset(13)};
  Greater greater_kernel{ctx};
  Tensor x = Tensor::FromInt64("", {4}, {-2, 0, 3, 7});
  Tensor y = Tensor::FromInt64("", {4}, {-1, 0, 1, 9});
  Tensor z = greater_kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  EXPECT_EQ(z.data_type, static_cast<int32_t>(core::runtime::DataType::BOOL));
  EXPECT_EQ(z.data[0], 0); // -2 > -1
  EXPECT_EQ(z.data[1], 0); //  0 >  0
  EXPECT_EQ(z.data[2], 1); //  3 >  1
  EXPECT_EQ(z.data[3], 0); //  7 >  9
}

TEST(KernelClass, LessClassMatchesReference) {
  const KernelContext ctx{DefaultOpset(13)};
  Less less_kernel{ctx};
  Tensor x = Tensor::FromFloat("", {2, 2}, {1.0f, 2.0f, 3.0f, 4.0f});
  Tensor y = Tensor::FromFloat("", {2, 2}, {2.0f, 2.0f, 2.0f, 2.0f});
  Tensor z = less_kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  EXPECT_EQ(z.data_type, static_cast<int32_t>(core::runtime::DataType::BOOL));
  EXPECT_EQ(z.data[0], 1); // 1 < 2 -> true
  EXPECT_EQ(z.data[1], 0); // 2 < 2 -> false
  EXPECT_EQ(z.data[2], 0); // 3 < 2 -> false
  EXPECT_EQ(z.data[3], 0); // 4 < 2 -> false
}

TEST(KernelClass, LessClassBroadcastsScalar) {
  const KernelContext ctx{DefaultOpset(13)};
  Less less_kernel{ctx};
  Tensor x = Tensor::FromFloat("", {2, 2}, {1.0f, 2.0f, 3.0f, 4.0f});
  Tensor y = Tensor::FromFloat("", {}, {2.5f});
  Tensor z = less_kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  EXPECT_EQ(z.data[0], 1);
  EXPECT_EQ(z.data[1], 1);
  EXPECT_EQ(z.data[2], 0);
  EXPECT_EQ(z.data[3], 0);
}

TEST(KernelClass, LessInPlaceWritesToPreallocatedOutput) {
  const KernelContext ctx{DefaultOpset(13)};
  Less less_kernel{ctx};
  Tensor x = Tensor::FromFloat("", {2, 2}, {1.0f, 2.0f, 3.0f, 4.0f});
  Tensor y = Tensor::FromFloat("", {2, 2}, {2.0f, 2.0f, 2.0f, 2.0f});
  Tensor z("", core::runtime::DataType::BOOL, {2, 2}, std::vector<uint8_t>(4, 9));
  less_kernel(x, y, z);
  EXPECT_EQ(z.data[0], 1);
  EXPECT_EQ(z.data[1], 0);
  EXPECT_EQ(z.data[2], 0);
  EXPECT_EQ(z.data[3], 0);
}

TEST(KernelClass, LessRejectsUnsupportedDtype) {
  // BOOL inputs are not in the supported dtype set (FLOAT/INT8/INT16/UINT8/
  // UINT16/UINT32/UINT64) so the kernel must reject them.
  const KernelContext ctx{DefaultOpset(13)};
  Less less_kernel{ctx};
  Tensor x("", core::runtime::DataType::BOOL, {2}, {1, 0});
  Tensor y("", core::runtime::DataType::BOOL, {2}, {0, 1});
  EXPECT_THROW((void)less_kernel(x, y), std::invalid_argument);
}

TEST(KernelClass, LessClassMatchesReferenceInt16) {
  const KernelContext ctx{DefaultOpset(13)};
  Less less_kernel{ctx};
  Tensor x = Tensor::FromInt16("", {4}, {-2, 0, 3, 7});
  Tensor y = Tensor::FromInt16("", {4}, {-1, 0, 1, 9});
  Tensor z = less_kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  EXPECT_EQ(z.data[0], 1); // -2 < -1
  EXPECT_EQ(z.data[1], 0); //  0 <  0
  EXPECT_EQ(z.data[2], 0); //  3 <  1
  EXPECT_EQ(z.data[3], 1); //  7 <  9
}

TEST(KernelClass, LessClassMatchesReferenceUint64) {
  const KernelContext ctx{DefaultOpset(13)};
  Less less_kernel{ctx};
  Tensor x = Tensor::FromUint64("", {4}, {1ull, 5ull, 0ull, 7ull});
  Tensor y = Tensor::FromUint64("", {4}, {2ull, 5ull, 0ull, 6ull});
  Tensor z = less_kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  EXPECT_EQ(z.data[0], 1);
  EXPECT_EQ(z.data[1], 0);
  EXPECT_EQ(z.data[2], 0);
  EXPECT_EQ(z.data[3], 0);
}

TEST(KernelClass, LessClassMatchesReferenceInt64) {
  const KernelContext ctx{DefaultOpset(13)};
  Less less_kernel{ctx};
  Tensor x = Tensor::FromInt64("", {4}, {-2, 0, 3, 7});
  Tensor y = Tensor::FromInt64("", {4}, {-1, 0, 1, 9});
  Tensor z = less_kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  EXPECT_EQ(z.data_type, static_cast<int32_t>(core::runtime::DataType::BOOL));
  EXPECT_EQ(z.data[0], 1); // -2 < -1
  EXPECT_EQ(z.data[1], 0); //  0 <  0
  EXPECT_EQ(z.data[2], 0); //  3 <  1
  EXPECT_EQ(z.data[3], 1); //  7 <  9
}

TEST(KernelClass, LessClassMatchesReferenceInt32) {
  const KernelContext ctx{DefaultOpset(13)};
  Less less_kernel{ctx};
  Tensor x = Tensor::FromInt32("", {4}, {-2, 0, 3, 7});
  Tensor y = Tensor::FromInt32("", {4}, {-1, 0, 1, 9});
  Tensor z = less_kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  EXPECT_EQ(z.data_type, static_cast<int32_t>(core::runtime::DataType::BOOL));
  EXPECT_EQ(z.data[0], 1);
  EXPECT_EQ(z.data[1], 0);
  EXPECT_EQ(z.data[2], 0);
  EXPECT_EQ(z.data[3], 1);
}

TEST(KernelClass, GreaterOrEqualClassMatchesReference) {
  const KernelContext ctx{DefaultOpset(16)};
  GreaterOrEqual ge_kernel{ctx};
  Tensor x = Tensor::FromFloat("", {2, 2}, {1.0f, 2.0f, 3.0f, 4.0f});
  Tensor y = Tensor::FromFloat("", {2, 2}, {2.0f, 2.0f, 2.0f, 2.0f});
  Tensor z = ge_kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  EXPECT_EQ(z.data_type, static_cast<int32_t>(core::runtime::DataType::BOOL));
  EXPECT_EQ(z.data[0], 0); // 1 >= 2 -> false
  EXPECT_EQ(z.data[1], 1); // 2 >= 2 -> true
  EXPECT_EQ(z.data[2], 1); // 3 >= 2 -> true
  EXPECT_EQ(z.data[3], 1); // 4 >= 2 -> true
}

TEST(KernelClass, GreaterOrEqualClassBroadcastsScalar) {
  const KernelContext ctx{DefaultOpset(16)};
  GreaterOrEqual ge_kernel{ctx};
  Tensor x = Tensor::FromFloat("", {2, 2}, {1.0f, 2.0f, 3.0f, 4.0f});
  Tensor y = Tensor::FromFloat("", {}, {2.0f});
  Tensor z = ge_kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  EXPECT_EQ(z.data[0], 0);
  EXPECT_EQ(z.data[1], 1);
  EXPECT_EQ(z.data[2], 1);
  EXPECT_EQ(z.data[3], 1);
}

TEST(KernelClass, GreaterOrEqualInPlaceWritesToPreallocatedOutput) {
  const KernelContext ctx{DefaultOpset(16)};
  GreaterOrEqual ge_kernel{ctx};
  Tensor x = Tensor::FromFloat("", {4}, {1.0f, 2.0f, 3.0f, 4.0f});
  Tensor y = Tensor::FromFloat("", {4}, {2.0f, 2.0f, 2.0f, 2.0f});
  Tensor output("", core::runtime::DataType::BOOL, {4}, {0, 0, 0, 0});
  ge_kernel(x, y, output);
  ASSERT_EQ(output.element_count(), 4);
  EXPECT_EQ(output.data[0], 0);
  EXPECT_EQ(output.data[1], 1);
  EXPECT_EQ(output.data[2], 1);
  EXPECT_EQ(output.data[3], 1);
}

TEST(KernelClass, GreaterOrEqualRejectsUnsupportedDtype) {
  const KernelContext ctx{DefaultOpset(16)};
  Tensor x("", core::runtime::DataType::BOOL, {2}, {1, 0});
  Tensor y("", core::runtime::DataType::BOOL, {2}, {1, 1});
  GreaterOrEqual ge_kernel{ctx};
  EXPECT_THROW({ (void)ge_kernel(x, y); }, std::invalid_argument);
}

TEST(KernelClass, LessOrEqualClassMatchesReference) {
  const KernelContext ctx{DefaultOpset(16)};
  LessOrEqual le_kernel{ctx};
  Tensor x = Tensor::FromFloat("", {2, 2}, {1.0f, 2.0f, 3.0f, 4.0f});
  Tensor y = Tensor::FromFloat("", {2, 2}, {2.0f, 2.0f, 2.0f, 2.0f});
  Tensor z = le_kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  EXPECT_EQ(z.data_type, static_cast<int32_t>(core::runtime::DataType::BOOL));
  EXPECT_EQ(z.data[0], 1); // 1 <= 2 -> true
  EXPECT_EQ(z.data[1], 1); // 2 <= 2 -> true
  EXPECT_EQ(z.data[2], 0); // 3 <= 2 -> false
  EXPECT_EQ(z.data[3], 0); // 4 <= 2 -> false
}

TEST(KernelClass, LessOrEqualClassBroadcastsScalar) {
  const KernelContext ctx{DefaultOpset(16)};
  LessOrEqual le_kernel{ctx};
  Tensor x = Tensor::FromFloat("", {2, 2}, {1.0f, 2.0f, 3.0f, 4.0f});
  Tensor y = Tensor::FromFloat("", {}, {2.0f});
  Tensor z = le_kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  EXPECT_EQ(z.data[0], 1);
  EXPECT_EQ(z.data[1], 1);
  EXPECT_EQ(z.data[2], 0);
  EXPECT_EQ(z.data[3], 0);
}

TEST(KernelClass, LessOrEqualInPlaceWritesToPreallocatedOutput) {
  const KernelContext ctx{DefaultOpset(16)};
  LessOrEqual le_kernel{ctx};
  Tensor x = Tensor::FromFloat("", {4}, {1.0f, 2.0f, 3.0f, 4.0f});
  Tensor y = Tensor::FromFloat("", {4}, {2.0f, 2.0f, 2.0f, 2.0f});
  Tensor output("", core::runtime::DataType::BOOL, {4}, {0, 0, 0, 0});
  le_kernel(x, y, output);
  ASSERT_EQ(output.element_count(), 4);
  EXPECT_EQ(output.data[0], 1);
  EXPECT_EQ(output.data[1], 1);
  EXPECT_EQ(output.data[2], 0);
  EXPECT_EQ(output.data[3], 0);
}

TEST(KernelClass, LessOrEqualRejectsUnsupportedDtype) {
  const KernelContext ctx{DefaultOpset(16)};
  Tensor x("", core::runtime::DataType::BOOL, {2}, {1, 0});
  Tensor y("", core::runtime::DataType::BOOL, {2}, {1, 1});
  LessOrEqual le_kernel{ctx};
  EXPECT_THROW({ (void)le_kernel(x, y); }, std::invalid_argument);
}

TEST(KernelClass, LessOrEqualClassMatchesReferenceInt64) {
  const KernelContext ctx{DefaultOpset(16)};
  LessOrEqual le_kernel{ctx};
  Tensor x = Tensor::FromInt64("", {4}, {-2, 0, 3, 7});
  Tensor y = Tensor::FromInt64("", {4}, {-1, 0, 1, 9});
  Tensor z = le_kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  EXPECT_EQ(z.data_type, static_cast<int32_t>(core::runtime::DataType::BOOL));
  EXPECT_EQ(z.data[0], 1); // -2 <= -1
  EXPECT_EQ(z.data[1], 1); //  0 <=  0
  EXPECT_EQ(z.data[2], 0); //  3 <=  1
  EXPECT_EQ(z.data[3], 1); //  7 <=  9
}

TEST(KernelClass, EqualClassMatchesReference) {
  const KernelContext ctx{DefaultOpset(19)};
  Equal equal_kernel{ctx};
  Tensor x = Tensor::FromInt32("", {2, 2}, {1, 2, 3, 4});
  Tensor y = Tensor::FromInt32("", {2, 2}, {1, 0, 3, 0});
  Tensor z = equal_kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  EXPECT_EQ(z.data_type, static_cast<int32_t>(core::runtime::DataType::BOOL));
  EXPECT_EQ(z.data[0], 1); // 1 == 1
  EXPECT_EQ(z.data[1], 0); // 2 == 0
  EXPECT_EQ(z.data[2], 1); // 3 == 3
  EXPECT_EQ(z.data[3], 0); // 4 == 0
}

TEST(KernelClass, EqualClassBroadcastsScalar) {
  const KernelContext ctx{DefaultOpset(19)};
  Equal equal_kernel{ctx};
  Tensor x = Tensor::FromInt32("", {2, 2}, {1, 2, 3, 4});
  Tensor y = Tensor::FromInt32("", {}, {3});
  Tensor z = equal_kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  EXPECT_EQ(z.data[0], 0);
  EXPECT_EQ(z.data[1], 0);
  EXPECT_EQ(z.data[2], 1);
  EXPECT_EQ(z.data[3], 0);
}

TEST(KernelClass, EqualInPlaceWritesToPreallocatedOutput) {
  const KernelContext ctx{DefaultOpset(19)};
  Equal equal_kernel{ctx};
  Tensor x = Tensor::FromInt32("", {3}, {1, 2, 3});
  Tensor y = Tensor::FromInt32("", {3}, {1, 0, 3});
  Tensor out("", core::runtime::DataType::BOOL, {3}, std::vector<uint8_t>(3));
  equal_kernel(x, y, out);
  EXPECT_EQ(out.data[0], 1);
  EXPECT_EQ(out.data[1], 0);
  EXPECT_EQ(out.data[2], 1);
}

TEST(KernelClass, EqualClassMatchesReferenceBool) {
  const KernelContext ctx{DefaultOpset(19)};
  Equal equal_kernel{ctx};
  Tensor x("", core::runtime::DataType::BOOL, {4}, {1, 0, 1, 0});
  Tensor y("", core::runtime::DataType::BOOL, {4}, {1, 1, 0, 0});
  Tensor z = equal_kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  EXPECT_EQ(z.data[0], 1);
  EXPECT_EQ(z.data[1], 0);
  EXPECT_EQ(z.data[2], 0);
  EXPECT_EQ(z.data[3], 1);
}

TEST(KernelClass, EqualClassMatchesReferenceFloat) {
  const KernelContext ctx{DefaultOpset(19)};
  Equal equal_kernel{ctx};
  Tensor x = Tensor::FromFloat("", {3}, {1.0f, 2.0f, 3.0f});
  Tensor y = Tensor::FromFloat("", {3}, {1.0f, 0.0f, 3.0f});
  Tensor z = equal_kernel(x, y);
  ASSERT_EQ(z.element_count(), 3);
  EXPECT_EQ(z.data[0], 1);
  EXPECT_EQ(z.data[1], 0);
  EXPECT_EQ(z.data[2], 1);
}

TEST(KernelClass, EqualClassMatchesReferenceString) {
  const KernelContext ctx{DefaultOpset(19)};
  Equal equal_kernel{ctx};
  Tensor x = Tensor::FromStrings("", {2}, {"string1", "string2"});
  Tensor y = Tensor::FromStrings("", {2}, {"string1", "string3"});
  Tensor z = equal_kernel(x, y);
  ASSERT_EQ(z.element_count(), 2);
  EXPECT_EQ(z.data_type, static_cast<int32_t>(core::runtime::DataType::BOOL));
  EXPECT_EQ(z.data[0], 1);
  EXPECT_EQ(z.data[1], 0);
}

TEST(KernelClass, EqualClassBroadcastsScalarString) {
  const KernelContext ctx{DefaultOpset(19)};
  Equal equal_kernel{ctx};
  Tensor x = Tensor::FromStrings("", {2}, {"string1", "string2"});
  Tensor y = Tensor::FromStrings("", {1}, {"string1"});
  Tensor z = equal_kernel(x, y);
  ASSERT_EQ(z.element_count(), 2);
  EXPECT_EQ(z.data[0], 1);
  EXPECT_EQ(z.data[1], 0);
}

TEST(KernelClass, EqualRejectsUnsupportedDtype) {
  const KernelContext ctx{DefaultOpset(19)};
  Equal equal_kernel{ctx};
  Tensor x("", core::runtime::DataType::COMPLEX64, {2}, std::vector<uint8_t>(16));
  Tensor y("", core::runtime::DataType::COMPLEX64, {2}, std::vector<uint8_t>(16));
  EXPECT_THROW((void)equal_kernel(x, y), std::invalid_argument);
}

TEST(KernelClass, WhereClassSelectsValuesElementwise) {
  const KernelContext ctx{DefaultOpset(16)};
  Where where_kernel{ctx};
  Tensor condition = Tensor::FromBool("condition", {2, 2}, {1, 0, 1, 0});
  Tensor x = Tensor::FromFloat("x", {2, 2}, {1.0f, 2.0f, 3.0f, 4.0f});
  Tensor y = Tensor::FromFloat("y", {2, 2}, {5.0f, 6.0f, 7.0f, 8.0f});
  Tensor output = where_kernel(condition, x, y);
  ASSERT_EQ(output.data_type, static_cast<int32_t>(core::runtime::DataType::FLOAT));
  ASSERT_EQ(output.shape, (std::vector<int64_t>{2, 2}));
  const float *values = output.AsFloat();
  EXPECT_FLOAT_EQ(values[0], 1.0f);
  EXPECT_FLOAT_EQ(values[1], 6.0f);
  EXPECT_FLOAT_EQ(values[2], 3.0f);
  EXPECT_FLOAT_EQ(values[3], 8.0f);
}

TEST(KernelClass, WhereClassBroadcastsInputs) {
  const KernelContext ctx{DefaultOpset(16)};
  Where where_kernel{ctx};
  Tensor condition = Tensor::FromBool("condition", {2, 1}, {1, 0});
  Tensor x = Tensor::FromInt32("x", {2, 3}, {1, 2, 3, 4, 5, 6});
  Tensor y = Tensor::FromInt32("y", {1, 3}, {10, 20, 30});
  Tensor output = where_kernel(condition, x, y);
  ASSERT_EQ(output.shape, (std::vector<int64_t>{2, 3}));
  const int32_t *values = output.AsInt32();
  EXPECT_EQ(values[0], 1);
  EXPECT_EQ(values[1], 2);
  EXPECT_EQ(values[2], 3);
  EXPECT_EQ(values[3], 10);
  EXPECT_EQ(values[4], 20);
  EXPECT_EQ(values[5], 30);
}

TEST(KernelClass, WhereRejectsNonBoolCondition) {
  const KernelContext ctx{DefaultOpset(16)};
  Where where_kernel{ctx};
  Tensor condition = Tensor::FromFloat("condition", {2}, {1.0f, 0.0f});
  Tensor x = Tensor::FromInt32("x", {2}, {1, 2});
  Tensor y = Tensor::FromInt32("y", {2}, {3, 4});
  EXPECT_THROW((void)where_kernel(condition, x, y), std::invalid_argument);
}

TEST(KernelClass, WhereUsesAllocatorWhenRuntimeContextHasOne) {
  const KernelContext ctx{DefaultOpset(16)};
  Where where_kernel{ctx};
  constexpr size_t kMaxAllocations = 1;
  SimpleRawBufferAllocator alloc(kMaxAllocations);
  RuntimeContext rt(core::runtime::RuntimeContextOptions{.allocator = &alloc});
  Tensor condition = Tensor::FromBool("condition", {2, 2}, {1, 0, 1, 0});
  Tensor x = Tensor::FromFloat("x", {2, 2}, {1.0f, 2.0f, 3.0f, 4.0f});
  Tensor y = Tensor::FromFloat("y", {2, 2}, {5.0f, 6.0f, 7.0f, 8.0f});

  Tensor output = where_kernel(condition, x, y, &rt);

  ASSERT_TRUE(output.has_allocation());
  ASSERT_NE(output.allocation(), nullptr);
  EXPECT_EQ(alloc.allocated_count(), 1u);
  EXPECT_EQ(output.bytes(), output.allocation()->data());
  EXPECT_EQ(output.data.size(), 0u);
  ASSERT_EQ(output.shape, (std::vector<int64_t>{2, 2}));
  const float *values = output.AsFloat();
  EXPECT_FLOAT_EQ(values[0], 1.0f);
  EXPECT_FLOAT_EQ(values[1], 6.0f);
  EXPECT_FLOAT_EQ(values[2], 3.0f);
  EXPECT_FLOAT_EQ(values[3], 8.0f);
}

// ---------------------------------------------------------------------------
// Bitwise kernels
// ---------------------------------------------------------------------------

TEST(KernelClass, BitwiseAndClassMatchesReferenceInt32) {
  const KernelContext ctx{DefaultOpset(18)};
  BitwiseAnd kernel{ctx};
  Tensor x = Tensor::FromInt32("", {4}, {0xF0, 0x0F, 0xAA, 0x55});
  Tensor y = Tensor::FromInt32("", {4}, {0xFF, 0xFF, 0x0F, 0xF0});
  Tensor z = kernel(x, y);
  ASSERT_EQ(z.data_type, static_cast<int32_t>(core::runtime::DataType::INT32));
  ASSERT_EQ(z.element_count(), 4);
  const int32_t *p = reinterpret_cast<const int32_t *>(z.data.data());
  EXPECT_EQ(p[0], 0xF0 & 0xFF);
  EXPECT_EQ(p[1], 0x0F & 0xFF);
  EXPECT_EQ(p[2], 0xAA & 0x0F);
  EXPECT_EQ(p[3], 0x55 & 0xF0);
}

TEST(KernelClass, BitwiseOrClassMatchesReferenceInt32) {
  const KernelContext ctx{DefaultOpset(18)};
  BitwiseOr kernel{ctx};
  Tensor x = Tensor::FromInt32("", {3}, {1, 2, 4});
  Tensor y = Tensor::FromInt32("", {3}, {8, 8, 8});
  Tensor z = kernel(x, y);
  const int32_t *p = reinterpret_cast<const int32_t *>(z.data.data());
  EXPECT_EQ(p[0], 9);
  EXPECT_EQ(p[1], 10);
  EXPECT_EQ(p[2], 12);
}

TEST(KernelClass, BitwiseXorClassMatchesReferenceInt32) {
  const KernelContext ctx{DefaultOpset(18)};
  BitwiseXor kernel{ctx};
  Tensor x = Tensor::FromInt32("", {3}, {0xFF, 0xAA, 0x0F});
  Tensor y = Tensor::FromInt32("", {3}, {0x0F, 0xFF, 0xF0});
  Tensor z = kernel(x, y);
  const int32_t *p = reinterpret_cast<const int32_t *>(z.data.data());
  EXPECT_EQ(p[0], 0xFF ^ 0x0F);
  EXPECT_EQ(p[1], 0xAA ^ 0xFF);
  EXPECT_EQ(p[2], 0x0F ^ 0xF0);
}

TEST(KernelClass, BitwiseAndClassBroadcastsScalar) {
  const KernelContext ctx{DefaultOpset(18)};
  BitwiseAnd kernel{ctx};
  Tensor x = Tensor::FromInt32("", {2, 2}, {0xF0, 0x0F, 0xAA, 0x55});
  Tensor y = Tensor::FromInt32("", {}, {0x0F});
  Tensor z = kernel(x, y);
  ASSERT_EQ(z.element_count(), 4);
  const int32_t *p = reinterpret_cast<const int32_t *>(z.data.data());
  EXPECT_EQ(p[0], 0x00);
  EXPECT_EQ(p[1], 0x0F);
  EXPECT_EQ(p[2], 0x0A);
  EXPECT_EQ(p[3], 0x05);
}

TEST(KernelClass, BitwiseAndInPlaceWritesToPreallocatedOutput) {
  const KernelContext ctx{DefaultOpset(18)};
  BitwiseAnd kernel{ctx};
  Tensor x = Tensor::FromInt32("", {3}, {0xF0, 0x0F, 0xAA});
  Tensor y = Tensor::FromInt32("", {3}, {0xFF, 0xFF, 0x0F});
  Tensor z("", core::runtime::DataType::INT32, {3}, std::vector<uint8_t>(3 * sizeof(int32_t)));
  kernel(x, y, z);
  const int32_t *p = reinterpret_cast<const int32_t *>(z.data.data());
  EXPECT_EQ(p[0], 0xF0);
  EXPECT_EQ(p[1], 0x0F);
  EXPECT_EQ(p[2], 0x0A);
}

TEST(KernelClass, BitwiseAndRejectsFloatTensors) {
  const KernelContext ctx{DefaultOpset(18)};
  BitwiseAnd kernel{ctx};
  Tensor x = Tensor::FromFloat("", {2}, {1.0f, 0.0f});
  Tensor y = Tensor::FromFloat("", {2}, {1.0f, 1.0f});
  EXPECT_THROW((void)kernel(x, y), std::invalid_argument);
}

TEST(KernelClass, BitwiseAndRejectsBoolTensors) {
  const KernelContext ctx{DefaultOpset(18)};
  BitwiseAnd kernel{ctx};
  Tensor x("", core::runtime::DataType::BOOL, {2}, {1, 0});
  Tensor y("", core::runtime::DataType::BOOL, {2}, {1, 1});
  EXPECT_THROW((void)kernel(x, y), std::invalid_argument);
}

TEST(KernelClass, BitwiseNotClassMatchesReferenceInt32) {
  const KernelContext ctx{DefaultOpset(18)};
  BitwiseNot kernel{ctx};
  Tensor x = Tensor::FromInt32("", {4}, {0, -1, 0xFF, 0x0F});
  Tensor z = kernel(x);
  ASSERT_EQ(z.data_type, static_cast<int32_t>(core::runtime::DataType::INT32));
  ASSERT_EQ(z.element_count(), 4);
  const int32_t *p = reinterpret_cast<const int32_t *>(z.data.data());
  EXPECT_EQ(p[0], ~0);
  EXPECT_EQ(p[1], ~(-1));
  EXPECT_EQ(p[2], ~0xFF);
  EXPECT_EQ(p[3], ~0x0F);
}

TEST(KernelClass, BitwiseNotClassMatchesReferenceUint8) {
  const KernelContext ctx{DefaultOpset(18)};
  BitwiseNot kernel{ctx};
  Tensor x = Tensor::FromUint8("", {3}, {0x00, 0xF0, 0xAA});
  Tensor z = kernel(x);
  ASSERT_EQ(z.data_type, static_cast<int32_t>(core::runtime::DataType::UINT8));
  EXPECT_EQ(z.data[0], static_cast<uint8_t>(0xFF));
  EXPECT_EQ(z.data[1], static_cast<uint8_t>(0x0F));
  EXPECT_EQ(z.data[2], static_cast<uint8_t>(0x55));
}

TEST(KernelClass, BitwiseNotInPlaceWritesToPreallocatedOutput) {
  const KernelContext ctx{DefaultOpset(18)};
  BitwiseNot kernel{ctx};
  Tensor x = Tensor::FromInt32("", {2}, {0, 0xFF});
  Tensor z("", core::runtime::DataType::INT32, {2}, std::vector<uint8_t>(2 * sizeof(int32_t)));
  kernel(x, z);
  const int32_t *p = reinterpret_cast<const int32_t *>(z.data.data());
  EXPECT_EQ(p[0], ~0);
  EXPECT_EQ(p[1], ~0xFF);
}

TEST(KernelClass, BitwiseNotRejectsBoolTensors) {
  const KernelContext ctx{DefaultOpset(18)};
  BitwiseNot kernel{ctx};
  Tensor x("", core::runtime::DataType::BOOL, {2}, {1, 0});
  EXPECT_THROW((void)kernel(x), std::invalid_argument);
}

TEST(KernelClass, NotClassMatchesReference) {
  const KernelContext ctx{DefaultOpset(1)};
  Not not_kernel{ctx};
  Tensor x("", core::runtime::DataType::BOOL, {4}, {1, 0, 1, 0});
  Tensor y = not_kernel(x);
  ASSERT_EQ(y.element_count(), 4);
  EXPECT_EQ(y.data_type, static_cast<int32_t>(core::runtime::DataType::BOOL));
  EXPECT_EQ(y.data[0], 0u);
  EXPECT_EQ(y.data[1], 1u);
  EXPECT_EQ(y.data[2], 0u);
  EXPECT_EQ(y.data[3], 1u);
}

TEST(KernelClass, NotUsesTypedParallelTuning) {
  const KernelContext ctx{DefaultOpset(1)};
  Not not_kernel{ctx};
  const auto bool_key = not_kernel.TuningKey(static_cast<int32_t>(core::runtime::DataType::BOOL));

  EXPECT_EQ(bool_key.library, "onnx_light");
  EXPECT_EQ(bool_key.kernel, "Not");
  EXPECT_EQ(bool_key.implementation, "portable");
  EXPECT_EQ(not_kernel.TuningKey(static_cast<int32_t>(core::runtime::DataType::FLOAT)).device,
            core::symbolic::Device::kUndefined);

  const auto schema = core::runtime::GetKernelTuningRegistry().FindSchema(bool_key);
  ASSERT_NE(schema, nullptr);
  EXPECT_EQ(schema->portable_defaults().Get<int64_t>("parallel.minimum_elements"),
            core::runtime::kParallelForGrainSize);

  core::runtime::KernelTuningParameters tuned{bool_key,
                                              {{"parallel.minimum_elements", int64_t{3}}}};
  not_kernel.Configure(tuned);
  EXPECT_EQ(not_kernel.tuning().parallel_minimum_elements, 3);

  Tensor x("", core::runtime::DataType::BOOL, {4}, {1, 0, 1, 0});
  Tensor y = not_kernel(x);
  EXPECT_EQ(y.data, (std::vector<uint8_t>{0, 1, 0, 1}));
}

TEST(KernelClass, NotInPlaceWritesToPreallocatedOutput) {
  const KernelContext ctx{DefaultOpset(1)};
  Not not_kernel{ctx};
  Tensor x("", core::runtime::DataType::BOOL, {3}, {0, 1, 0});
  Tensor y("", core::runtime::DataType::BOOL, {3}, std::vector<uint8_t>(3));
  not_kernel(x, y);
  EXPECT_EQ(y.data[0], 1u);
  EXPECT_EQ(y.data[1], 0u);
  EXPECT_EQ(y.data[2], 1u);
}

TEST(KernelClass, NotRejectsNonBoolTensors) {
  const KernelContext ctx{DefaultOpset(1)};
  Not not_kernel{ctx};
  Tensor x = Tensor::FromInt32("", {2}, {0, 1});
  EXPECT_THROW((void)not_kernel(x), std::invalid_argument);
}

TEST(KernelClass, IsNaNClassMatchesReference) {
  const KernelContext ctx{DefaultOpset(20)};
  IsNaN isnan_kernel{ctx};
  const float nan_v = std::numeric_limits<float>::quiet_NaN();
  const float inf_v = std::numeric_limits<float>::infinity();
  Tensor x = Tensor::FromFloat("", {6}, {-1.2f, nan_v, inf_v, 2.8f, -inf_v, inf_v});
  Tensor y = isnan_kernel(x);
  ASSERT_EQ(y.element_count(), 6);
  EXPECT_EQ(y.data_type, static_cast<int32_t>(core::runtime::DataType::BOOL));
  EXPECT_EQ(y.data[0], 0u);
  EXPECT_EQ(y.data[1], 1u);
  EXPECT_EQ(y.data[2], 0u);
  EXPECT_EQ(y.data[3], 0u);
  EXPECT_EQ(y.data[4], 0u);
  EXPECT_EQ(y.data[5], 0u);
}

TEST(KernelClass, IsNaNRejectsNonFloatTensors) {
  const KernelContext ctx{DefaultOpset(20)};
  IsNaN isnan_kernel{ctx};
  Tensor x = Tensor::FromInt32("", {2}, {0, 1});
  EXPECT_THROW((void)isnan_kernel(x), std::invalid_argument);
}

TEST(KernelClass, IsNaNDoubleMatchesReference) {
  const KernelContext ctx{DefaultOpset(20)};
  IsNaN isnan_kernel{ctx};
  const double nan_v = std::numeric_limits<double>::quiet_NaN();
  const double inf_v = std::numeric_limits<double>::infinity();
  Tensor x = Tensor::FromDouble("", {6}, {-1.2, nan_v, inf_v, 2.8, -inf_v, inf_v});
  Tensor y = isnan_kernel(x);
  ASSERT_EQ(y.element_count(), 6);
  EXPECT_EQ(y.data_type, static_cast<int32_t>(core::runtime::DataType::BOOL));
  EXPECT_EQ(y.data[0], 0u);
  EXPECT_EQ(y.data[1], 1u);
  EXPECT_EQ(y.data[2], 0u);
  EXPECT_EQ(y.data[3], 0u);
  EXPECT_EQ(y.data[4], 0u);
  EXPECT_EQ(y.data[5], 0u);
}

TEST(KernelClass, IsNaNFloat16MatchesReference) {
  const KernelContext ctx{DefaultOpset(20)};
  IsNaN isnan_kernel{ctx};
  const float nan_v = std::numeric_limits<float>::quiet_NaN();
  const float inf_v = std::numeric_limits<float>::infinity();
  Tensor x = onnx_kernels::MakeFloat16Tensor("", {6}, {-1.2f, nan_v, inf_v, 2.8f, -inf_v, inf_v});
  Tensor y = isnan_kernel(x);
  ASSERT_EQ(y.element_count(), 6);
  EXPECT_EQ(y.data_type, static_cast<int32_t>(core::runtime::DataType::BOOL));
  EXPECT_EQ(y.data[0], 0u);
  EXPECT_EQ(y.data[1], 1u);
  EXPECT_EQ(y.data[2], 0u);
  EXPECT_EQ(y.data[3], 0u);
  EXPECT_EQ(y.data[4], 0u);
  EXPECT_EQ(y.data[5], 0u);
}

TEST(KernelClass, IsNaNBfloat16MatchesReference) {
  const KernelContext ctx{DefaultOpset(20)};
  IsNaN isnan_kernel{ctx};
  const float nan_v = std::numeric_limits<float>::quiet_NaN();
  const float inf_v = std::numeric_limits<float>::infinity();
  const std::vector<float> vals = {-1.2f, nan_v, inf_v, 2.8f, -inf_v, inf_v};
  std::vector<uint8_t> raw(vals.size() * sizeof(uint16_t));
  auto *dst = reinterpret_cast<uint16_t *>(raw.data());
  for (size_t i = 0; i < vals.size(); ++i)
    dst[i] = onnx_kernels::FloatToBfloat16Bits(vals[i]);
  Tensor x("", static_cast<int32_t>(core::runtime::DataType::BFLOAT16), {6}, std::move(raw));
  Tensor y = isnan_kernel(x);
  ASSERT_EQ(y.element_count(), 6);
  EXPECT_EQ(y.data_type, static_cast<int32_t>(core::runtime::DataType::BOOL));
  EXPECT_EQ(y.data[0], 0u);
  EXPECT_EQ(y.data[1], 1u);
  EXPECT_EQ(y.data[2], 0u);
  EXPECT_EQ(y.data[3], 0u);
  EXPECT_EQ(y.data[4], 0u);
  EXPECT_EQ(y.data[5], 0u);
}

TEST(KernelClass, IsInfClassMatchesReference) {
  const KernelContext ctx{DefaultOpset(20)};
  IsInf isinf_kernel{ctx};
  const float nan_v = std::numeric_limits<float>::quiet_NaN();
  const float inf_v = std::numeric_limits<float>::infinity();
  Tensor x = Tensor::FromFloat("", {6}, {-1.2f, nan_v, inf_v, 2.8f, -inf_v, inf_v});

  Tensor both = isinf_kernel(x);
  ASSERT_EQ(both.element_count(), 6);
  EXPECT_EQ(both.data_type, static_cast<int32_t>(core::runtime::DataType::BOOL));
  EXPECT_EQ(both.data[0], 0u);
  EXPECT_EQ(both.data[1], 0u);
  EXPECT_EQ(both.data[2], 1u);
  EXPECT_EQ(both.data[3], 0u);
  EXPECT_EQ(both.data[4], 1u);
  EXPECT_EQ(both.data[5], 1u);

  Tensor only_pos = isinf_kernel(x, /*detect_positive=*/1, /*detect_negative=*/0);
  EXPECT_EQ(only_pos.data[2], 1u);
  EXPECT_EQ(only_pos.data[4], 0u);
  EXPECT_EQ(only_pos.data[5], 1u);

  Tensor only_neg = isinf_kernel(x, /*detect_positive=*/0, /*detect_negative=*/1);
  EXPECT_EQ(only_neg.data[2], 0u);
  EXPECT_EQ(only_neg.data[4], 1u);
  EXPECT_EQ(only_neg.data[5], 0u);
}

TEST(KernelClass, IsInfRejectsNonFloatTensors) {
  const KernelContext ctx{DefaultOpset(20)};
  IsInf isinf_kernel{ctx};
  Tensor x = Tensor::FromInt32("", {2}, {0, 1});
  EXPECT_THROW((void)isinf_kernel(x), std::invalid_argument);
}

} // namespace Test
