Add validation and unit tests

Author: yuslepukhin

onnxruntime/core/providers/cpu/tensor/affine_grid.cc

@@ -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"
 
@@ -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);
@@ -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);
@@ -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).colwise() + 2x1).transpose()
 }
 
 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;
@@ -113,9 +115,15 @@ 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], "]");
+
     TensorShape grid_shape{N, H, W, 2};
     auto grid = context->Output(0, grid_shape);
 
@@ -128,9 +136,15 @@ 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], "]");
+
     TensorShape grid_shape{N, D, H, W, 3};
     auto grid = context->Output(0, grid_shape);
 

onnxruntime/test/providers/cpu/tensor/affine_grid_test.cc

@@ -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"  // Ensure this include is present
 
 namespace onnxruntime {
 namespace test {
@@ -178,5 +179,95 @@ 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. It must be 4 for 2D or 5 for 5D");
+}
+
+// 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. It must be 4 for 2D or 5 for 5D");
+}
+
+// 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");
+}
 }  // namespace test
 }  // namespace onnxruntime