Add validation and unit tests AffineGrid

Merged
YU
8mo ago
+220-20

Description

This pull request improves the robustness and correctness of the AffineGrid operator in ONNX Runtime by adding stricter input validation and updating the implementation for better type safety. It also introduces comprehensive unit tests to verify error handling for invalid inputs.

Input validation enhancements

  • Added explicit checks in affine_grid.cc to ensure that the batch dimension in size matches the batch dimension in theta, and that theta has the correct shape for both 2D ([N, 2, 3]) and 3D ([N, 3, 4]) cases. This prevents silent errors due to shape mismatches. [1] [2]

Type safety and code clarity improvements

  • Updated grid generation functions to use the template type T for Eigen vector creation, replacing hardcoded float types. This ensures correct behavior for different data types and improves code clarity. [1] [2]
  • Added clarifying comments to matrix and array initializations in affine grid generator functions to make the code easier to understand and maintain.

Test coverage improvements

  • Added a series of unit tests in affine_grid_test.cc to validate input shape checks, including cases for incorrect tensor dimensions, mismatched batch sizes, and invalid size lengths. These tests ensure that the operator fails gracefully with informative error messages when given invalid inputs.
  • Ensured the required test utility header is included to support new tests.

Code cleanup

  • Removed unused includes from affine_grid.cc to tidy up the codebase.