// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

// Custom GridSample tests not generated by grid_sample_test_gen.py.
// These cover edge cases (extreme coordinates, NaN/Inf, dim==1, border, 3D)
// that exercise the float->int64 cast hardening in
// onnxruntime/core/providers/cpu/tensor/grid_sample.cc.

#include <limits>
#include <memory>
#include <utility>
#include <vector>

#include "gtest/gtest.h"

#include "core/framework/execution_provider.h"
#include "test/providers/provider_test_utils.h"
#include "test/util/include/default_providers.h"

namespace onnxruntime {
namespace test {
namespace {

std::vector<std::unique_ptr<IExecutionProvider>> GetCpuExecutionProviders() {
  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
  execution_providers.emplace_back(DefaultCpuExecutionProvider());
  return execution_providers;
}

std::vector<std::unique_ptr<IExecutionProvider>> GetExecutionProviders() {
  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;

  execution_providers.emplace_back(DefaultCpuExecutionProvider());

#ifdef USE_CUDA
  execution_providers.emplace_back(DefaultCudaExecutionProvider());
#ifdef ENABLE_CUDA_NHWC_OPS
  execution_providers.push_back(DefaultCudaNHWCExecutionProvider());
#endif
#endif

#if defined(USE_COREML)
  execution_providers.push_back(DefaultCoreMLExecutionProvider(/*use_mlprogram*/ true));
#endif

#ifdef USE_WEBGPU
  execution_providers.push_back(DefaultWebGpuExecutionProvider());
#endif

  return execution_providers;
}

template <typename T>
void RunTests(T& test, std::vector<std::unique_ptr<IExecutionProvider>>&& execution_providers) {
  for (size_t idx = 0; idx < execution_providers.size(); ++idx) {
    test.ConfigEp(std::move(execution_providers[idx])).RunWithConfig();
  }
  execution_providers.clear();
}

template <typename T>
void RunCpuOnly(T& test) {
  RunTests(test, GetCpuExecutionProviders());
}

}  // namespace

template <typename T>
class GridSampleCustomTest : public ::testing::Test {
};

using GridSampleCustomTestTypes = ::testing::Types<float, MLFloat16>;
TYPED_TEST_SUITE(GridSampleCustomTest, GridSampleCustomTestTypes);

TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_linear_zeros_mixed_bounds_right_bottom) {
  // Crafts grid points that mix fully in-bounds sampling with cases where either the right, bottom,
  // or both neighbors fall outside the source image so zero padding must be applied. This ensures
  // the optimized bilinear fast path matches the generic implementation for boundary handling.
  OpTester test("GridSample", 20);
  std::string mode = "linear";
  std::string padding_mode = "zeros";
  int64_t align_corners = 0;
  std::initializer_list<int64_t> X_shape{1, 1, 2, 2};
  std::initializer_list<TypeParam> X_data{TypeParam(1.0f), TypeParam(2.0f), TypeParam(3.0f), TypeParam(4.0f)};
  std::initializer_list<int64_t> Grid_shape{1, 2, 2, 2};
  // (nx, ny) pairs: center (in-bounds), right edge (x out), bottom edge (y out), corner (both out)
  std::initializer_list<TypeParam> Grid_data{
      TypeParam(0.0f), TypeParam(0.0f),   // center (all neighbors in bounds)
      TypeParam(0.9f), TypeParam(0.0f),   // near right edge (right neighbors out of bounds)
      TypeParam(0.0f), TypeParam(0.9f),   // near bottom edge (bottom neighbors out)
      TypeParam(0.9f), TypeParam(0.9f)};  // near bottom-right corner (both right and bottom neighbors out)
  std::initializer_list<int64_t> Y_shape{1, 1, 2, 2};
  std::initializer_list<TypeParam> Y_data{
      TypeParam(2.5f),    // all neighbors in bounds
      TypeParam(1.8f),    // right neighbors partially out-of-bounds
      TypeParam(2.1f),    // bottom neighbors partially out-of-bounds
      TypeParam(1.44f)};  // both right and bottom neighbors out-of-bounds
  test.AddInput<TypeParam>("X", X_shape, X_data);
  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
  test.AddAttribute("mode", mode);
  test.AddAttribute("padding_mode", padding_mode);
  test.AddAttribute("align_corners", align_corners);
  test.AddOutput<TypeParam>("Y", Y_shape, Y_data);
  RunTests(test, GetExecutionProviders());
}

TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_linear_zeros_mixed_bounds_left_top) {
  // Similar to test_grid_sample_20_4D_linear_zeros_mixed_bounds_right_bottom but focuses on left/top boundary cases,
  // where the left and/or top neighbors fall outside the source image and zero padding must be applied.
  // This ensures the optimized bilinear fast path correctly handles left/top boundary conditions.
  OpTester test("GridSample", 20);
  std::string mode = "linear";
  std::string padding_mode = "zeros";
  int64_t align_corners = 0;
  std::initializer_list<int64_t> X_shape{1, 1, 2, 2};
  std::initializer_list<TypeParam> X_data{TypeParam(1.0f), TypeParam(2.0f), TypeParam(3.0f), TypeParam(4.0f)};
  std::initializer_list<int64_t> Grid_shape{1, 2, 2, 2};
  // (nx, ny) pairs: center (in-bounds), left edge (x out), top edge (y out), corner (both out)
  std::initializer_list<TypeParam> Grid_data{
      TypeParam(0.0f), TypeParam(0.0f),     // center (all neighbors in bounds)
      TypeParam(-0.9f), TypeParam(0.0f),    // near left edge (left neighbors out of bounds)
      TypeParam(0.0f), TypeParam(-0.9f),    // near top edge (top neighbors out of bounds)
      TypeParam(-0.9f), TypeParam(-0.9f)};  // near top-left corner (both left and top neighbors out of bounds)
  std::initializer_list<int64_t> Y_shape{1, 1, 2, 2};
  std::initializer_list<TypeParam> Y_data{
      TypeParam(2.5f),    // all neighbors in bounds
      TypeParam(1.2f),    // left neighbors partially out-of-bounds
      TypeParam(0.9f),    // top neighbors partially out-of-bounds
      TypeParam(0.36f)};  // both left and top neighbors out-of-bounds
  test.AddInput<TypeParam>("X", X_shape, X_data);
  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
  test.AddAttribute("mode", mode);
  test.AddAttribute("padding_mode", padding_mode);
  test.AddAttribute("align_corners", align_corners);
  test.AddOutput<TypeParam>("Y", Y_shape, Y_data);
  RunTests(test, GetExecutionProviders());
}

TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_nearest_reflection_extreme_coords) {
  // Regression test: large grid coordinates previously produced wild indices in the
  // reflection branch of PixelAtGrid. The internal division-and-int-cast inside
  // GsReflect now uses int64_t, and PixelAtGrid clamps the resulting index back into
  // the image range, so sampling stays in bounds. Using "nearest" mode keeps the test
  // focused on the reflection index path (no bilinear weight arithmetic), and using a
  // constant-valued input makes the expected output deterministic regardless of which
  // in-range pixel is hit.
  OpTester test("GridSample", 20);
  test.AddAttribute("mode", std::string("nearest"));
  test.AddAttribute("padding_mode", std::string("reflection"));
  test.AddAttribute("align_corners", int64_t{0});

  std::initializer_list<int64_t> X_shape{1, 1, 2, 2};
  std::initializer_list<TypeParam> X_data{TypeParam(1.0f), TypeParam(1.0f),
                                          TypeParam(1.0f), TypeParam(1.0f)};

  std::initializer_list<int64_t> Grid_shape{1, 1, 2, 2};
  // 1e+10 exceeds INT_MAX, so the pre-fix `int n = static_cast<int>(dx / range)`
  // inside GsReflect would invoke undefined behavior; 1e+10 still fits comfortably
  // in int64_t. For MLFloat16 the value overflows to +/-Inf, exercising the same
  // safety-clamp fallback in PixelAtGrid.
  std::initializer_list<TypeParam> Grid_data{
      TypeParam(1.0e+10f), TypeParam(1.0e+10f),
      TypeParam(-1.0e+10f), TypeParam(-1.0e+10f)};

  std::initializer_list<int64_t> Y_shape{1, 1, 1, 2};

  test.AddInput<TypeParam>("X", X_shape, X_data);
  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
  test.AddOutput<TypeParam>("Y", Y_shape, {TypeParam(1.0f), TypeParam(1.0f)});
  RunCpuOnly(test);
}

TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_nearest_reflection_nan_coords) {
  // Regression test: NaN grid coordinates would otherwise propagate into
  // static_cast<int64_t>(std::nearbyint(NaN)), which is undefined behavior.
  // The Compute loop now sanitizes such coordinates to the in-range border
  // before performing the cast, producing defined sampling.
  OpTester test("GridSample", 20);
  test.AddAttribute("mode", std::string("nearest"));
  test.AddAttribute("padding_mode", std::string("reflection"));
  test.AddAttribute("align_corners", int64_t{0});

  std::initializer_list<int64_t> X_shape{1, 1, 2, 2};
  std::initializer_list<TypeParam> X_data{TypeParam(1.0f), TypeParam(1.0f),
                                          TypeParam(1.0f), TypeParam(1.0f)};

  std::initializer_list<int64_t> Grid_shape{1, 1, 2, 2};
  const float nan_val = std::numeric_limits<float>::quiet_NaN();
  std::initializer_list<TypeParam> Grid_data{
      TypeParam(nan_val), TypeParam(nan_val),
      TypeParam(0.0f), TypeParam(nan_val)};

  std::initializer_list<int64_t> Y_shape{1, 1, 1, 2};

  test.AddInput<TypeParam>("X", X_shape, X_data);
  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
  // With a constant-valued image any reflected pixel returns 1.0.
  test.AddOutput<TypeParam>("Y", Y_shape, {TypeParam(1.0f), TypeParam(1.0f)});
  RunCpuOnly(test);
}

TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_bilinear_reflection_infinity_coords) {
  // Regression test: Inf grid coordinates would otherwise reach
  // static_cast<int64_t>(std::floor(Inf)), which is undefined behavior.
  OpTester test("GridSample", 20);
  test.AddAttribute("mode", std::string("linear"));
  test.AddAttribute("padding_mode", std::string("reflection"));
  test.AddAttribute("align_corners", int64_t{0});

  std::initializer_list<int64_t> X_shape{1, 1, 2, 2};
  std::initializer_list<TypeParam> X_data{TypeParam(1.0f), TypeParam(1.0f),
                                          TypeParam(1.0f), TypeParam(1.0f)};

  std::initializer_list<int64_t> Grid_shape{1, 1, 2, 2};
  const float inf_val = std::numeric_limits<float>::infinity();
  std::initializer_list<TypeParam> Grid_data{
      TypeParam(inf_val), TypeParam(-inf_val),
      TypeParam(-inf_val), TypeParam(inf_val)};

  std::initializer_list<int64_t> Y_shape{1, 1, 1, 2};

  test.AddInput<TypeParam>("X", X_shape, X_data);
  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
  // Constant-valued image: any reflected sample returns 1.0.
  test.AddOutput<TypeParam>("Y", Y_shape, {TypeParam(1.0f), TypeParam(1.0f)});
  RunCpuOnly(test);
}

TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_bilinear_reflection_extreme_coords) {
  // Bilinear interpolation path: the four corner indices are computed via
  // static_cast<int64_t>(std::floor(x)). With a constant-valued image the
  // expected output is 1.0 regardless of which reflected pixel is hit.
  OpTester test("GridSample", 20);
  test.AddAttribute("mode", std::string("linear"));
  test.AddAttribute("padding_mode", std::string("reflection"));
  test.AddAttribute("align_corners", int64_t{0});

  std::initializer_list<int64_t> X_shape{1, 1, 2, 2};
  std::initializer_list<TypeParam> X_data{TypeParam(1.0f), TypeParam(1.0f),
                                          TypeParam(1.0f), TypeParam(1.0f)};

  std::initializer_list<int64_t> Grid_shape{1, 1, 2, 2};
  std::initializer_list<TypeParam> Grid_data{
      TypeParam(1.0e+20f), TypeParam(1.0e+20f),
      TypeParam(-1.0e+20f), TypeParam(-1.0e+20f)};

  std::initializer_list<int64_t> Y_shape{1, 1, 1, 2};

  test.AddInput<TypeParam>("X", X_shape, X_data);
  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
  // With a constant-valued image any sampled pixel is 1.0.
  test.AddOutput<TypeParam>("Y", Y_shape, {TypeParam(1.0f), TypeParam(1.0f)});
  RunCpuOnly(test);
}

TEST(GridSampleCustomTestFloatOnly, test_grid_sample_20_4D_cubic_reflection_extreme_coords) {
  // Cubic interpolation path: the 4x4 neighborhood is computed via
  // static_cast<int64_t>(std::floor(x)) - 1, exercising the same UB-prone cast.
  // Cubic only supports float (T1 type constraint), so this test is float-only.
  OpTester test("GridSample", 20);
  test.AddAttribute("mode", std::string("cubic"));
  test.AddAttribute("padding_mode", std::string("reflection"));
  test.AddAttribute("align_corners", int64_t{0});

  std::initializer_list<int64_t> X_shape{1, 1, 4, 4};
  std::initializer_list<float> X_data{
      1.0f, 1.0f, 1.0f, 1.0f,
      1.0f, 1.0f, 1.0f, 1.0f,
      1.0f, 1.0f, 1.0f, 1.0f,
      1.0f, 1.0f, 1.0f, 1.0f};

  std::initializer_list<int64_t> Grid_shape{1, 1, 1, 2};
  std::initializer_list<float> Grid_data{1.0e+10f, -1.0e+10f};

  std::initializer_list<int64_t> Y_shape{1, 1, 1, 1};

  test.AddInput<float>("X", X_shape, X_data);
  test.AddInput<float>("Grid", Grid_shape, Grid_data);
  // Constant-valued image: cubic interpolation of all-1.0 neighborhood is 1.0.
  test.AddOutput<float>("Y", Y_shape, {1.0f});
  RunCpuOnly(test);
}

TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_5D_nearest_reflection_extreme_coords) {
  // 3D / 5-D input: the reflection branch in PixelAtGrid3D received the same
  // safety clamp as PixelAtGrid; this test exercises that path.
  OpTester test("GridSample", 20);
  test.AddAttribute("mode", std::string("nearest"));
  test.AddAttribute("padding_mode", std::string("reflection"));
  test.AddAttribute("align_corners", int64_t{0});

  std::initializer_list<int64_t> X_shape{1, 1, 2, 2, 2};
  std::initializer_list<TypeParam> X_data{
      TypeParam(1.0f), TypeParam(1.0f), TypeParam(1.0f), TypeParam(1.0f),
      TypeParam(1.0f), TypeParam(1.0f), TypeParam(1.0f), TypeParam(1.0f)};

  std::initializer_list<int64_t> Grid_shape{1, 1, 1, 2, 3};
  std::initializer_list<TypeParam> Grid_data{
      TypeParam(1.0e+10f), TypeParam(1.0e+10f), TypeParam(1.0e+10f),
      TypeParam(-1.0e+10f), TypeParam(-1.0e+10f), TypeParam(-1.0e+10f)};

  std::initializer_list<int64_t> Y_shape{1, 1, 1, 1, 2};

  test.AddInput<TypeParam>("X", X_shape, X_data);
  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
  test.AddOutput<TypeParam>("Y", Y_shape, {TypeParam(1.0f), TypeParam(1.0f)});
  RunCpuOnly(test);
}

TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_nearest_reflection_dim1_align_corners) {
  // With a 1x1 image and align_corners=true, x_min == x_max so the reflection
  // range is zero. GsReflect must not divide by zero on any out-of-range coord.
  OpTester test("GridSample", 20);
  test.AddAttribute("mode", std::string("nearest"));
  test.AddAttribute("padding_mode", std::string("reflection"));
  test.AddAttribute("align_corners", int64_t{1});

  std::initializer_list<int64_t> X_shape{1, 1, 1, 1};
  std::initializer_list<TypeParam> X_data{TypeParam(7.0f)};

  std::initializer_list<int64_t> Grid_shape{1, 1, 2, 2};
  std::initializer_list<TypeParam> Grid_data{
      TypeParam(0.5f), TypeParam(0.5f),
      TypeParam(-0.5f), TypeParam(-0.5f)};

  std::initializer_list<int64_t> Y_shape{1, 1, 1, 2};

  test.AddInput<TypeParam>("X", X_shape, X_data);
  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
  // The single input pixel is the only valid sample for any reflected coordinate.
  test.AddOutput<TypeParam>("Y", Y_shape, {TypeParam(7.0f), TypeParam(7.0f)});
  RunCpuOnly(test);
}

TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_bilinear_border_extreme_coords) {
  // Border padding mode is independently safe via std::clamp; this regression
  // test ensures it stays safe with extreme grid coordinates.
  OpTester test("GridSample", 20);
  test.AddAttribute("mode", std::string("linear"));
  test.AddAttribute("padding_mode", std::string("border"));
  test.AddAttribute("align_corners", int64_t{0});

  std::initializer_list<int64_t> X_shape{1, 1, 2, 2};
  std::initializer_list<TypeParam> X_data{TypeParam(1.0f), TypeParam(1.0f),
                                          TypeParam(1.0f), TypeParam(1.0f)};

  std::initializer_list<int64_t> Grid_shape{1, 1, 2, 2};
  std::initializer_list<TypeParam> Grid_data{
      TypeParam(1.0e+20f), TypeParam(1.0e+20f),
      TypeParam(-1.0e+20f), TypeParam(-1.0e+20f)};

  std::initializer_list<int64_t> Y_shape{1, 1, 1, 2};

  test.AddInput<TypeParam>("X", X_shape, X_data);
  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
  // Constant-valued image: any clamped border sample is 1.0.
  test.AddOutput<TypeParam>("Y", Y_shape, {TypeParam(1.0f), TypeParam(1.0f)});
  RunCpuOnly(test);
}

TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_linear_zeros_extreme_nan_inf_coords) {
  // Zeros padding is safe by design, but the IsSafeForInt64Conversion sanitization
  // (in PrecomputeBilinearSamplePlan2D) replaces NaN/Inf/extreme values with the
  // lower border (-0.5) before the floor cast. Verify the optimized bilinear
  // zeros-padding fast path still produces a deterministic, well-defined output
  // for those inputs.
  OpTester test("GridSample", 20);
  test.AddAttribute("mode", std::string("linear"));
  test.AddAttribute("padding_mode", std::string("zeros"));
  test.AddAttribute("align_corners", int64_t{0});

  std::initializer_list<int64_t> X_shape{1, 1, 2, 2};
  std::initializer_list<TypeParam> X_data{TypeParam(1.0f), TypeParam(1.0f),
                                          TypeParam(1.0f), TypeParam(1.0f)};

  std::initializer_list<int64_t> Grid_shape{1, 1, 3, 2};
  std::initializer_list<TypeParam> Grid_data{
      TypeParam(std::numeric_limits<float>::quiet_NaN()),
      TypeParam(std::numeric_limits<float>::quiet_NaN()),
      TypeParam(std::numeric_limits<float>::infinity()),
      TypeParam(-std::numeric_limits<float>::infinity()),
      TypeParam(1.0e+20f), TypeParam(-1.0e+20f)};

  std::initializer_list<int64_t> Y_shape{1, 1, 1, 3};

  test.AddInput<TypeParam>("X", X_shape, X_data);
  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
  // Each unsafe coord is sanitized to (-0.5, -0.5); only the bottom-right neighbor
  // (pixel(0, 0) = 1.0) is in-bounds with weight 0.5*0.5 = 0.25, others are zeros.
  test.AddOutput<TypeParam>("Y", Y_shape, {TypeParam(0.25f), TypeParam(0.25f), TypeParam(0.25f)});
  RunCpuOnly(test);
}

TEST(GridSampleCustomTestFloatOnly, test_grid_sample_20_4D_cubic_reflection_nan_inf_coords) {
  // Cubic interpolation also goes through static_cast<int64_t>(std::floor(x)) - 1.
  // Cover the NaN / +Inf / -Inf paths through the cubic kernel; constant-valued
  // image makes the output trivially 1.0 regardless of which 4x4 neighborhood is
  // selected after sanitization.
  OpTester test("GridSample", 20);
  test.AddAttribute("mode", std::string("cubic"));
  test.AddAttribute("padding_mode", std::string("reflection"));
  test.AddAttribute("align_corners", int64_t{0});

  std::initializer_list<int64_t> X_shape{1, 1, 4, 4};
  std::initializer_list<float> X_data{
      1.0f, 1.0f, 1.0f, 1.0f,
      1.0f, 1.0f, 1.0f, 1.0f,
      1.0f, 1.0f, 1.0f, 1.0f,
      1.0f, 1.0f, 1.0f, 1.0f};

  std::initializer_list<int64_t> Grid_shape{1, 1, 3, 2};
  std::initializer_list<float> Grid_data{
      std::numeric_limits<float>::quiet_NaN(),
      std::numeric_limits<float>::quiet_NaN(),
      std::numeric_limits<float>::infinity(),
      -std::numeric_limits<float>::infinity(),
      -std::numeric_limits<float>::infinity(),
      std::numeric_limits<float>::infinity()};

  std::initializer_list<int64_t> Y_shape{1, 1, 1, 3};

  test.AddInput<float>("X", X_shape, X_data);
  test.AddInput<float>("Grid", Grid_shape, Grid_data);
  test.AddOutput<float>("Y", Y_shape, {1.0f, 1.0f, 1.0f});
  RunCpuOnly(test);
}

TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_5D_linear_reflection_extreme_coords) {
  // 3D trilinear path: exercises std::floor cast for x, y, and z with extreme
  // finite coordinates. Sanitization replaces the unsafe components with the
  // lower borders before any int64 cast.
  OpTester test("GridSample", 20);
  test.AddAttribute("mode", std::string("linear"));
  test.AddAttribute("padding_mode", std::string("reflection"));
  test.AddAttribute("align_corners", int64_t{0});

  std::initializer_list<int64_t> X_shape{1, 1, 2, 2, 2};
  std::initializer_list<TypeParam> X_data{
      TypeParam(1.0f), TypeParam(1.0f), TypeParam(1.0f), TypeParam(1.0f),
      TypeParam(1.0f), TypeParam(1.0f), TypeParam(1.0f), TypeParam(1.0f)};

  std::initializer_list<int64_t> Grid_shape{1, 1, 1, 2, 3};
  std::initializer_list<TypeParam> Grid_data{
      TypeParam(1.0e+20f), TypeParam(1.0e+20f), TypeParam(1.0e+20f),
      TypeParam(-1.0e+20f), TypeParam(-1.0e+20f), TypeParam(-1.0e+20f)};

  std::initializer_list<int64_t> Y_shape{1, 1, 1, 1, 2};

  test.AddInput<TypeParam>("X", X_shape, X_data);
  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
  // Constant-valued image: trilinear interpolation of all-1.0 neighborhood is 1.0.
  test.AddOutput<TypeParam>("Y", Y_shape, {TypeParam(1.0f), TypeParam(1.0f)});
  RunCpuOnly(test);
}

TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_5D_nearest_reflection_nan_inf_coords) {
  // 3D path with NaN / Inf / -Inf in different spatial dimensions. Verifies
  // that the 3D sanitization branch handles non-finite values along any axis.
  OpTester test("GridSample", 20);
  test.AddAttribute("mode", std::string("nearest"));
  test.AddAttribute("padding_mode", std::string("reflection"));
  test.AddAttribute("align_corners", int64_t{0});

  std::initializer_list<int64_t> X_shape{1, 1, 2, 2, 2};
  std::initializer_list<TypeParam> X_data{
      TypeParam(1.0f), TypeParam(1.0f), TypeParam(1.0f), TypeParam(1.0f),
      TypeParam(1.0f), TypeParam(1.0f), TypeParam(1.0f), TypeParam(1.0f)};

  std::initializer_list<int64_t> Grid_shape{1, 1, 1, 2, 3};
  std::initializer_list<TypeParam> Grid_data{
      TypeParam(std::numeric_limits<float>::quiet_NaN()),
      TypeParam(std::numeric_limits<float>::infinity()),
      TypeParam(-std::numeric_limits<float>::infinity()),
      TypeParam(std::numeric_limits<float>::infinity()),
      TypeParam(std::numeric_limits<float>::quiet_NaN()),
      TypeParam(-std::numeric_limits<float>::infinity())};

  std::initializer_list<int64_t> Y_shape{1, 1, 1, 1, 2};

  test.AddInput<TypeParam>("X", X_shape, X_data);
  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
  test.AddOutput<TypeParam>("Y", Y_shape, {TypeParam(1.0f), TypeParam(1.0f)});
  RunCpuOnly(test);
}

TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_nearest_reflection_align_corners_extreme) {
  // align_corners=1 on a non-degenerate (3x3) image. With align_corners=1 the
  // lower border becomes 0 (not -0.5), so unsafe coords are sanitized to (0, 0)
  // and resolve to the top-left pixel of the image.
  OpTester test("GridSample", 20);
  test.AddAttribute("mode", std::string("nearest"));
  test.AddAttribute("padding_mode", std::string("reflection"));
  test.AddAttribute("align_corners", int64_t{1});

  std::initializer_list<int64_t> X_shape{1, 1, 3, 3};
  std::initializer_list<TypeParam> X_data{
      TypeParam(1.0f), TypeParam(2.0f), TypeParam(3.0f),
      TypeParam(4.0f), TypeParam(5.0f), TypeParam(6.0f),
      TypeParam(7.0f), TypeParam(8.0f), TypeParam(9.0f)};

  std::initializer_list<int64_t> Grid_shape{1, 1, 2, 2};
  std::initializer_list<TypeParam> Grid_data{
      TypeParam(1.0e+20f), TypeParam(1.0e+20f),
      TypeParam(std::numeric_limits<float>::quiet_NaN()),
      TypeParam(-std::numeric_limits<float>::infinity())};

  std::initializer_list<int64_t> Y_shape{1, 1, 1, 2};

  test.AddInput<TypeParam>("X", X_shape, X_data);
  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
  // Both unsafe coords sanitize to (0, 0) -> pixel(0, 0) = 1.0.
  test.AddOutput<TypeParam>("Y", Y_shape, {TypeParam(1.0f), TypeParam(1.0f)});
  RunCpuOnly(test);
}

TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_nearest_reflection_negative_only_extreme) {
  // Asymmetry: all extreme coordinates are negative. Sanitization replaces them
  // with x_min/y_min = -0.5, and nearbyint(-0.5) rounds to 0 (banker's rounding),
  // so each output samples pixel(0, 0) of a non-constant image.
  OpTester test("GridSample", 20);
  test.AddAttribute("mode", std::string("nearest"));
  test.AddAttribute("padding_mode", std::string("reflection"));
  test.AddAttribute("align_corners", int64_t{0});

  std::initializer_list<int64_t> X_shape{1, 1, 3, 3};
  std::initializer_list<TypeParam> X_data{
      TypeParam(1.0f), TypeParam(2.0f), TypeParam(3.0f),
      TypeParam(4.0f), TypeParam(5.0f), TypeParam(6.0f),
      TypeParam(7.0f), TypeParam(8.0f), TypeParam(9.0f)};

  std::initializer_list<int64_t> Grid_shape{1, 1, 2, 2};
  std::initializer_list<TypeParam> Grid_data{
      TypeParam(-1.0e+20f), TypeParam(-1.0e+20f),
      TypeParam(-std::numeric_limits<float>::infinity()),
      TypeParam(-std::numeric_limits<float>::infinity())};

  std::initializer_list<int64_t> Y_shape{1, 1, 1, 2};

  test.AddInput<TypeParam>("X", X_shape, X_data);
  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
  // Sanitized to (-0.5, -0.5); nearbyint -> (0, 0); pixel(0, 0) = 1.0.
  test.AddOutput<TypeParam>("Y", Y_shape, {TypeParam(1.0f), TypeParam(1.0f)});
  RunCpuOnly(test);
}

TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_nearest_reflection_mixed_normal_pathological) {
  // Mixed normal and pathological grid coordinates in the same call. Validates
  // that the sanitization only replaces unsafe values and leaves normal grid
  // points untouched, sampling them correctly from a non-constant image.
  OpTester test("GridSample", 20);
  test.AddAttribute("mode", std::string("nearest"));
  test.AddAttribute("padding_mode", std::string("reflection"));
  test.AddAttribute("align_corners", int64_t{1});

  std::initializer_list<int64_t> X_shape{1, 1, 3, 3};
  std::initializer_list<TypeParam> X_data{
      TypeParam(1.0f), TypeParam(2.0f), TypeParam(3.0f),
      TypeParam(4.0f), TypeParam(5.0f), TypeParam(6.0f),
      TypeParam(7.0f), TypeParam(8.0f), TypeParam(9.0f)};

  // align_corners=1 on a 3x3 image maps normalized [-1, 1] -> pixel coords [0, 2].
  // Mix four points: center, top-left, NaN (sanitized -> (0,0)), and asymmetric
  // extreme (sanitized -> (0,0)).
  std::initializer_list<int64_t> Grid_shape{1, 1, 4, 2};
  std::initializer_list<TypeParam> Grid_data{
      TypeParam(0.0f), TypeParam(0.0f),    // -> pixel(1, 1) = 5
      TypeParam(-1.0f), TypeParam(-1.0f),  // -> pixel(0, 0) = 1
      TypeParam(std::numeric_limits<float>::quiet_NaN()),
      TypeParam(std::numeric_limits<float>::quiet_NaN()),  // sanitized -> pixel(0, 0) = 1
      TypeParam(1.0e+20f), TypeParam(-1.0e+20f)};          // sanitized -> pixel(0, 0) = 1

  std::initializer_list<int64_t> Y_shape{1, 1, 1, 4};

  test.AddInput<TypeParam>("X", X_shape, X_data);
  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
  test.AddOutput<TypeParam>("Y", Y_shape,
                            {TypeParam(5.0f), TypeParam(1.0f), TypeParam(1.0f), TypeParam(1.0f)});
  RunCpuOnly(test);
}

}  // namespace test
}  // namespace onnxruntime
