Skip to content

Add validation and unit tests AffineGrid #27572

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yuslepukhin merged 7 commits into from
Mar 10, 2026
Merged

Add validation and unit tests AffineGrid #27572

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
55385a4
Add validation and unit tests
yuslepukhin Mar 5, 2026
1816e95
Update onnxruntime/core/providers/cpu/tensor/affine_grid.cc
yuslepukhin Mar 5, 2026
2ed48db
Update onnxruntime/test/providers/cpu/tensor/affine_grid_test.cc
yuslepukhin Mar 5, 2026
7b31b06
Shorten substrings
yuslepukhin Mar 5, 2026
dcc9b57
Add positivity checks for H,W,D
yuslepukhin Mar 6, 2026
6e7f328
Disable a test that timesout
yuslepukhin Mar 9, 2026
767de79
Merge branch 'main' into yuslepukhin/affine_grid_validate
yuslepukhin Mar 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 38 additions & 19 deletions onnxruntime/core/providers/cpu/tensor/affine_grid.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,6 @@
#include "core/common/common.h"
#include "core/providers/op_kernel_type_control.h"
#include "core/util/math_cpuonly.h"
#include <iostream>
#include "Eigen/src/Core/Map.h"
#include <Eigen/Dense>
#include "core/common/eigen_common_wrapper.h"

Expand All @@ -28,13 +26,14 @@ REGISTER_KERNEL_TYPED(double)

template <typename T>
void generate_base_grid_2d(int64_t H, int64_t W, bool align_corners, Eigen::Matrix<T, Eigen::Dynamic, 2>& base_grid) {
Eigen::VectorXf row_vec = Eigen::VectorXf::LinSpaced(static_cast<Eigen::Index>(W), -1, 1);
using VectorT = Eigen::Matrix<T, Eigen::Dynamic, 1>;
VectorT row_vec = VectorT::LinSpaced(static_cast<Eigen::Index>(W), static_cast<T>(-1), static_cast<T>(1));
if (!align_corners) {
row_vec = row_vec * (W - 1) / W;
row_vec = row_vec * static_cast<T>(W - 1) / static_cast<T>(W);
}
Eigen::VectorXf col_vec = Eigen::VectorXf::LinSpaced(static_cast<Eigen::Index>(H), -1, 1);
VectorT col_vec = VectorT::LinSpaced(static_cast<Eigen::Index>(H), static_cast<T>(-1), static_cast<T>(1));
if (!align_corners) {
col_vec = col_vec * (H - 1) / H;
col_vec = col_vec * static_cast<T>(H - 1) / static_cast<T>(H);
}

base_grid.resize(static_cast<Eigen::Index>(H * W), 2);
Expand All @@ -47,17 +46,18 @@ void generate_base_grid_2d(int64_t H, int64_t W, bool align_corners, Eigen::Matr

template <typename T>
void generate_base_grid_3d(int64_t D, int64_t H, int64_t W, bool align_corners, Eigen::Matrix<T, Eigen::Dynamic, 3>& base_grid) {
Eigen::VectorXf row_vec = Eigen::VectorXf::LinSpaced(static_cast<Eigen::Index>(W), -1, 1);
using VectorT = Eigen::Matrix<T, Eigen::Dynamic, 1>;
VectorT row_vec = VectorT::LinSpaced(static_cast<Eigen::Index>(W), static_cast<T>(-1), static_cast<T>(1));
if (!align_corners) {
row_vec = row_vec * (W - 1) / W;
row_vec = row_vec * static_cast<T>(W - 1) / static_cast<T>(W);
}
Eigen::VectorXf col_vec = Eigen::VectorXf::LinSpaced(static_cast<Eigen::Index>(H), -1, 1);
VectorT col_vec = VectorT::LinSpaced(static_cast<Eigen::Index>(H), static_cast<T>(-1), static_cast<T>(1));
if (!align_corners) {
col_vec = col_vec * (H - 1) / H;
col_vec = col_vec * static_cast<T>(H - 1) / static_cast<T>(H);
}
Eigen::VectorXf slice_vec = Eigen::VectorXf::LinSpaced(static_cast<Eigen::Index>(D), -1, 1);
VectorT slice_vec = VectorT::LinSpaced(static_cast<Eigen::Index>(D), static_cast<T>(-1), static_cast<T>(1));
if (!align_corners) {
slice_vec = slice_vec * (D - 1) / D;
slice_vec = slice_vec * static_cast<T>(D - 1) / static_cast<T>(D);
}

base_grid.resize(static_cast<Eigen::Index>(D * H * W), 3);
Expand All @@ -75,25 +75,27 @@ void affine_grid_generator_2d(const Tensor* theta, const Eigen::Matrix<T, 2, Eig
const Eigen::StorageOptions option = Eigen::RowMajor;
auto theta_batch_offset = batch_num * 2 * 3;
const T* theta_data = theta->Data<T>() + theta_batch_offset;
const Eigen::Matrix<T, 2, 2, option> theta_R{{theta_data[0], theta_data[1]}, {theta_data[3], theta_data[4]}};
const Eigen::Array<T, 2, 1> theta_T(theta_data[2], theta_data[5]);
const Eigen::Matrix<T, 2, 2, option> theta_R{{theta_data[0], theta_data[1]}, {theta_data[3], theta_data[4]}}; // 2x2
const Eigen::Array<T, 2, 1> theta_T(theta_data[2], theta_data[5]); // 2x1

auto grid_batch_offset = batch_num * H * W * 2;
T* grid_data = grid->MutableData<T>() + grid_batch_offset;
Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, 2, option>> grid_matrix(grid_data, narrow<size_t>(H * W), 2);
grid_matrix = ((theta_R * base_grid_transposed).array().colwise() + theta_T).matrix().transpose();
grid_matrix = ((theta_R * base_grid_transposed).array().colwise() + theta_T).matrix().transpose(); // ((2x2 * 2xN).array().colwise() + 2x1).matrix().transpose() => Nx2
}

template <typename T>
void affine_grid_generator_3d(const Tensor* theta, const Eigen::Matrix<T, 3, Eigen::Dynamic>& base_grid_transposed, int64_t batch_num, int64_t D, int64_t H, int64_t W, Tensor* grid) {
const Eigen::StorageOptions option = Eigen::RowMajor;
auto theta_batch_offset = batch_num * 3 * 4;
const T* theta_data = theta->Data<T>() + theta_batch_offset;

const Eigen::Matrix<T, 3, 3, option> theta_R{
{theta_data[0], theta_data[1], theta_data[2]},
{theta_data[4], theta_data[5], theta_data[6]},
{theta_data[8], theta_data[9], theta_data[10]}};
const Eigen::Array<T, 3, 1> theta_T(theta_data[3], theta_data[7], theta_data[11]);
{theta_data[8], theta_data[9], theta_data[10]}}; // 3x3

const Eigen::Array<T, 3, 1> theta_T(theta_data[3], theta_data[7], theta_data[11]); // 3x1

auto grid_batch_offset = batch_num * D * H * W * 3;
T* grid_data = grid->MutableData<T>() + grid_batch_offset;
Expand All @@ -113,9 +115,17 @@ Status AffineGrid<T>::Compute(OpKernelContext* context) const {
const TensorShape& size_shape = size->Shape();
const int64_t* size_data = size->Data<int64_t>();

if (size_shape.GetDims()[0] == 4 /*&& get_check_2d_grid_sample_consistency(theta_shape, size_shape, N, C, H, W)*/) {
if (size_shape.GetDims()[0] == 4) {
int64_t N = size_data[0], H = size_data[2], W = size_data[3];

ORT_RETURN_IF(N != theta_shape[0],
"AffineGrid: size[0] (", N, ") must equal theta batch dimension (", theta_shape[0], ")");
ORT_RETURN_IF(theta_shape[1] != 2 || theta_shape[2] != 3,
"AffineGrid: theta shape must be [N, 2, 3] for 2D, got [",
theta_shape[0], ", ", theta_shape[1], ", ", theta_shape[2], "]");
ORT_RETURN_IF(H <= 0, "AffineGrid: size[2] (H=", H, ") must be positive");
ORT_RETURN_IF(W <= 0, "AffineGrid: size[3] (W=", W, ") must be positive");

TensorShape grid_shape{N, H, W, 2};
auto grid = context->Output(0, grid_shape);

Expand All @@ -128,9 +138,18 @@ Status AffineGrid<T>::Compute(OpKernelContext* context) const {
};

concurrency::ThreadPool::TryBatchParallelFor(context->GetOperatorThreadPool(), narrow<size_t>(N), std::move(fn), 0);
} else if (size_shape.GetDims()[0] == 5 /*&& get_check_2d_grid_sample_consistency(theta_shape, size_shape, N, C, H, W)*/) {
} else if (size_shape.GetDims()[0] == 5) {
int64_t N = size_data[0], D = size_data[2], H = size_data[3], W = size_data[4];

ORT_RETURN_IF(N != theta_shape[0],
"AffineGrid: size[0] (", N, ") must equal theta batch dimension (", theta_shape[0], ")");
ORT_RETURN_IF(theta_shape[1] != 3 || theta_shape[2] != 4,
"AffineGrid: theta shape must be [N, 3, 4] for 3D, got [",
theta_shape[0], ", ", theta_shape[1], ", ", theta_shape[2], "]");
ORT_RETURN_IF(D <= 0, "AffineGrid: size[2] (D=", D, ") must be positive");
ORT_RETURN_IF(H <= 0, "AffineGrid: size[3] (H=", H, ") must be positive");
ORT_RETURN_IF(W <= 0, "AffineGrid: size[4] (W=", W, ") must be positive");

TensorShape grid_shape{N, D, H, W, 3};
auto grid = context->Output(0, grid_shape);

Expand Down
2 changes: 1 addition & 1 deletion onnxruntime/test/contrib_ops/matmul_bnb4_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -115,7 +115,7 @@ void RunTest(int64_t quant_type, int64_t M, int64_t N, int64_t K, int64_t block_
}
}

TEST(MatMulBnb4, Float32) {
TEST(MatMulBnb4, DISABLED_Float32) {
for (auto qt : {0, 1}) {
for (auto M : {1, 2, 100}) {
for (auto N : {1, 2, 32, 288}) {
Expand Down
181 changes: 181 additions & 0 deletions onnxruntime/test/providers/cpu/tensor/affine_grid_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
#include "core/util/math.h"
#include "gtest/gtest.h"
#include "test/providers/provider_test_utils.h"
#include "test/unittest_util/op_tester.h" // Provides OpTester used in these tests

namespace onnxruntime {
namespace test {
Expand Down Expand Up @@ -178,5 +179,185 @@ TEST(AffineGridTest, test_3d_7) {
test.SetOutputTolerance(0.0001f);
test.Run();
}

// Validation tests for input shape checks

// Test: theta must be a 3D tensor
TEST(AffineGridTest, invalid_theta_not_3d) {
OpTester test("AffineGrid", 20);
// theta is 2D instead of 3D
test.AddInput<float>("theta", {2, 6}, {1.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f});
test.AddInput<int64_t>("size", {4}, {2, 1, 2, 3});
test.AddOutput<float>("grid", {2, 2, 3, 2}, std::vector<float>(24, 0.0f));
test.Run(OpTester::ExpectResult::kExpectFailure, "AffineGrid : Input theta tensor dimension is not 3");
}

// Test: size length must be 4 or 5
TEST(AffineGridTest, invalid_size_length_3) {
OpTester test("AffineGrid", 20);
test.AddInput<float>("theta", {1, 2, 3}, {1.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f});
test.AddInput<int64_t>("size", {3}, {1, 1, 2});
test.AddOutput<float>("grid", {1, 2, 2}, std::vector<float>(4, 0.0f));
test.Run(OpTester::ExpectResult::kExpectFailure, "Length of input 'size' is 3");
}

// Test: size length must be 4 or 5 (too long)
TEST(AffineGridTest, invalid_size_length_6) {
OpTester test("AffineGrid", 20);
test.AddInput<float>("theta", {1, 2, 3}, {1.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f});
test.AddInput<int64_t>("size", {6}, {1, 1, 2, 3, 4, 5});
test.AddOutput<float>("grid", {1, 2, 3, 2}, std::vector<float>(12, 0.0f));
test.Run(OpTester::ExpectResult::kExpectFailure, "Length of input 'size' is 6");
}

// Test: 2D - batch dimension mismatch between theta and size
TEST(AffineGridTest, invalid_2d_batch_mismatch) {
OpTester test("AffineGrid", 20);
// theta has N=1, but size has N=2
test.AddInput<float>("theta", {1, 2, 3}, {1.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f});
test.AddInput<int64_t>("size", {4}, {2, 1, 2, 3});
test.AddOutput<float>("grid", {2, 2, 3, 2}, std::vector<float>(24, 0.0f));
test.Run(OpTester::ExpectResult::kExpectFailure, "must equal theta batch dimension");
}

// Test: 2D - theta shape must be [N, 2, 3], wrong second dimension
TEST(AffineGridTest, invalid_2d_theta_wrong_dim1) {
OpTester test("AffineGrid", 20);
// theta is [1, 3, 3] but for 2D it must be [N, 2, 3]
test.AddInput<float>("theta", {1, 3, 3}, {1.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f, 0.0f, 1.0f});
test.AddInput<int64_t>("size", {4}, {1, 1, 2, 3});
test.AddOutput<float>("grid", {1, 2, 3, 2}, std::vector<float>(12, 0.0f));
test.Run(OpTester::ExpectResult::kExpectFailure, "theta shape must be [N, 2, 3] for 2D");
}

// Test: 2D - theta shape must be [N, 2, 3], wrong third dimension
TEST(AffineGridTest, invalid_2d_theta_wrong_dim2) {
OpTester test("AffineGrid", 20);
// theta is [1, 2, 4] but for 2D it must be [N, 2, 3]
test.AddInput<float>("theta", {1, 2, 4}, {1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f});
test.AddInput<int64_t>("size", {4}, {1, 1, 2, 3});
test.AddOutput<float>("grid", {1, 2, 3, 2}, std::vector<float>(12, 0.0f));
test.Run(OpTester::ExpectResult::kExpectFailure, "theta shape must be [N, 2, 3] for 2D");
}

// Test: 3D - batch dimension mismatch between theta and size
TEST(AffineGridTest, invalid_3d_batch_mismatch) {
OpTester test("AffineGrid", 20);
// theta has N=1, but size has N=2
test.AddInput<float>("theta", {1, 3, 4}, {1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f});
test.AddInput<int64_t>("size", {5}, {2, 1, 2, 2, 3});
test.AddOutput<float>("grid", {2, 2, 2, 3, 3}, std::vector<float>(72, 0.0f));
test.Run(OpTester::ExpectResult::kExpectFailure, "must equal theta batch dimension");
}

// Test: 3D - theta shape must be [N, 3, 4], wrong second dimension
TEST(AffineGridTest, invalid_3d_theta_wrong_dim1) {
OpTester test("AffineGrid", 20);
// theta is [1, 2, 4] but for 3D it must be [N, 3, 4]
test.AddInput<float>("theta", {1, 2, 4}, {1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f});
test.AddInput<int64_t>("size", {5}, {1, 1, 2, 2, 3});
test.AddOutput<float>("grid", {1, 2, 2, 3, 3}, std::vector<float>(36, 0.0f));
test.Run(OpTester::ExpectResult::kExpectFailure, "theta shape must be [N, 3, 4] for 3D");
}

// Test: 3D - theta shape must be [N, 3, 4], wrong third dimension
TEST(AffineGridTest, invalid_3d_theta_wrong_dim2) {
OpTester test("AffineGrid", 20);
// theta is [1, 3, 3] but for 3D it must be [N, 3, 4]
test.AddInput<float>("theta", {1, 3, 3}, {1.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f, 0.0f, 1.0f});
test.AddInput<int64_t>("size", {5}, {1, 1, 2, 2, 3});
test.AddOutput<float>("grid", {1, 2, 2, 3, 3}, std::vector<float>(36, 0.0f));
test.Run(OpTester::ExpectResult::kExpectFailure, "theta shape must be [N, 3, 4] for 3D");
}

// Test: 2D - H must be positive (zero)
TEST(AffineGridTest, invalid_2d_H_zero) {
OpTester test("AffineGrid", 20);
test.AddInput<float>("theta", {1, 2, 3}, {1.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f});
test.AddInput<int64_t>("size", {4}, {1, 1, 0, 3});
test.AddOutput<float>("grid", {1, 0, 3, 2}, std::vector<float>(0, 0.0f));
test.Run(OpTester::ExpectResult::kExpectFailure, "size[2] (H=0) must be positive");
}

// Test: 2D - H must be positive (negative)
TEST(AffineGridTest, invalid_2d_H_negative) {
OpTester test("AffineGrid", 20);
test.AddInput<float>("theta", {1, 2, 3}, {1.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f});
test.AddInput<int64_t>("size", {4}, {1, 1, -1, 3});
test.AddOutput<float>("grid", {1, 1, 3, 2}, std::vector<float>(6, 0.0f));
test.Run(OpTester::ExpectResult::kExpectFailure, "size[2] (H=-1) must be positive");
}

// Test: 2D - W must be positive (zero)
TEST(AffineGridTest, invalid_2d_W_zero) {
OpTester test("AffineGrid", 20);
test.AddInput<float>("theta", {1, 2, 3}, {1.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f});
test.AddInput<int64_t>("size", {4}, {1, 1, 2, 0});
test.AddOutput<float>("grid", {1, 2, 0, 2}, std::vector<float>(0, 0.0f));
test.Run(OpTester::ExpectResult::kExpectFailure, "size[3] (W=0) must be positive");
}

// Test: 2D - W must be positive (negative)
TEST(AffineGridTest, invalid_2d_W_negative) {
OpTester test("AffineGrid", 20);
test.AddInput<float>("theta", {1, 2, 3}, {1.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f});
test.AddInput<int64_t>("size", {4}, {1, 1, 2, -2});
test.AddOutput<float>("grid", {1, 2, 1, 2}, std::vector<float>(4, 0.0f));
test.Run(OpTester::ExpectResult::kExpectFailure, "size[3] (W=-2) must be positive");
}

// Test: 3D - D must be positive (zero)
TEST(AffineGridTest, invalid_3d_D_zero) {
OpTester test("AffineGrid", 20);
test.AddInput<float>("theta", {1, 3, 4}, {1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f});
test.AddInput<int64_t>("size", {5}, {1, 1, 0, 2, 3});
test.AddOutput<float>("grid", {1, 0, 2, 3, 3}, std::vector<float>(0, 0.0f));
test.Run(OpTester::ExpectResult::kExpectFailure, "size[2] (D=0) must be positive");
}

// Test: 3D - D must be positive (negative)
TEST(AffineGridTest, invalid_3d_D_negative) {
OpTester test("AffineGrid", 20);
test.AddInput<float>("theta", {1, 3, 4}, {1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f});
test.AddInput<int64_t>("size", {5}, {1, 1, -1, 2, 3});
test.AddOutput<float>("grid", {1, 1, 2, 3, 3}, std::vector<float>(18, 0.0f));
test.Run(OpTester::ExpectResult::kExpectFailure, "size[2] (D=-1) must be positive");
}

// Test: 3D - H must be positive (zero)
TEST(AffineGridTest, invalid_3d_H_zero) {
OpTester test("AffineGrid", 20);
test.AddInput<float>("theta", {1, 3, 4}, {1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f});
test.AddInput<int64_t>("size", {5}, {1, 1, 2, 0, 3});
test.AddOutput<float>("grid", {1, 2, 0, 3, 3}, std::vector<float>(0, 0.0f));
test.Run(OpTester::ExpectResult::kExpectFailure, "size[3] (H=0) must be positive");
}

// Test: 3D - H must be positive (negative)
TEST(AffineGridTest, invalid_3d_H_negative) {
OpTester test("AffineGrid", 20);
test.AddInput<float>("theta", {1, 3, 4}, {1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f});
test.AddInput<int64_t>("size", {5}, {1, 1, 2, -3, 3});
test.AddOutput<float>("grid", {1, 2, 1, 3, 3}, std::vector<float>(18, 0.0f));
test.Run(OpTester::ExpectResult::kExpectFailure, "size[3] (H=-3) must be positive");
}

// Test: 3D - W must be positive (zero)
TEST(AffineGridTest, invalid_3d_W_zero) {
OpTester test("AffineGrid", 20);
test.AddInput<float>("theta", {1, 3, 4}, {1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f});
test.AddInput<int64_t>("size", {5}, {1, 1, 2, 3, 0});
test.AddOutput<float>("grid", {1, 2, 3, 0, 3}, std::vector<float>(0, 0.0f));
test.Run(OpTester::ExpectResult::kExpectFailure, "size[4] (W=0) must be positive");
}

// Test: 3D - W must be positive (negative)
TEST(AffineGridTest, invalid_3d_W_negative) {
OpTester test("AffineGrid", 20);
test.AddInput<float>("theta", {1, 3, 4}, {1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f});
test.AddInput<int64_t>("size", {5}, {1, 1, 2, 3, -4});
test.AddOutput<float>("grid", {1, 2, 3, 1, 3}, std::vector<float>(18, 0.0f));
test.Run(OpTester::ExpectResult::kExpectFailure, "size[4] (W=-4) must be positive");
}
} // namespace test
} // namespace onnxruntime
Loading