Skip to content

Validate grad_output in grid_sampler_{2d,3d}_backward #5584

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
torchxpubot merged 2 commits into from
Oct 2, 2026
Merged

Validate grad_output in grid_sampler_{2d,3d}_backward #5584

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a8b9dbb
Validate `grad_output` in `grid_sampler_{2d,3d}_backward`
pbielak Oct 1, 2026
7f2d044
Fix dispatch name
pbielak Oct 2, 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
14 changes: 9 additions & 5 deletions src/ATen/native/xpu/sycl/GridSampler.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -682,8 +682,7 @@ void grid_sampler_2d_backward_kernel(
int64_t padding_mode,
bool align_corners,
std::array<bool, 2> output_mask) {
check_grid_sampler_common(input, grid);
check_grid_sampler_2d(input, grid);
check_grid_sampler_2d_backward(input, grid, grad_output);

globalContext().alertNotDeterministic("grid_sampler_2d_backward_xpu");
auto N = input.size(0);
Expand Down Expand Up @@ -1491,8 +1490,13 @@ void grid_sampler_3d_backward_kernel(
std::array<bool, 2> output_mask) {
// See NOTE [ grid_sampler Native Functions ].
// Add checks here in case this is called instead of grid_sampler.
check_grid_sampler_common(input, grid);
check_grid_sampler_3d(input, grid, interpolation_mode);
check_grid_sampler_3d_backward(input, grid, grad_output);
TORCH_CHECK_NOT_IMPLEMENTED(
static_cast<GridSamplerInterpolation>(interpolation_mode) !=
GridSamplerInterpolation::Bicubic,
"grid_sampler(): bicubic interpolation with 5D input is not "
"implemented for ",
input.device().type());

globalContext().alertNotDeterministic("grid_sampler_3d_backward_xpu");
auto input_requires_grad = output_mask[0];
Expand All @@ -1506,7 +1510,7 @@ void grid_sampler_3d_backward_kernel(
at::ScalarType::BFloat16,
at::ScalarType::Half,
input.scalar_type(),
"grid_sampler_2d_backward_xpu",
"grid_sampler_3d_backward_xpu",
Comment thread
dvrogozh marked this conversation as resolved.
Resolved
[&] {
if (canUse32BitIndexMath(input) && canUse32BitIndexMath(grid) &&
canUse32BitIndexMath(grad_output)) {
Expand Down
31 changes: 31 additions & 0 deletions test/xpu/test_nn_xpu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14201,6 +14201,37 @@ def test_save_lstm_compatibility(self, device):
result = loaded_model(x)
self.assertEqual(result, expected)

def test_grid_sample_backward_error_checking(self, device):
Comment thread
dvrogozh marked this conversation as resolved.
Resolved
input = torch.empty(1, 1, 2, 2, device=device)
grid = torch.empty(1, 1, 1, 2, device=device)
grad_output = torch.empty(1, 1, 1, 1, device=device)

# assert no error
torch.ops.aten.grid_sampler_2d_backward(
grad_output, input, grid, 0, 0, False, (True, True)
)

with self.assertRaisesRegex(ValueError, "expected grad_output to have sizes"):
invalid_grad_output = torch.empty(2, 1, 1, 1, device=device)
torch.ops.aten.grid_sampler_2d_backward(
invalid_grad_output, input, grid, 0, 0, False, (True, True)
)

input = torch.empty(1, 1, 2, 2, 2, device=device)
grid = torch.empty(1, 1, 1, 1, 3, device=device)
grad_output = torch.empty(1, 1, 1, 1, 1, device=device)

# assert no error
torch.ops.aten.grid_sampler_3d_backward(
grad_output, input, grid, 0, 0, False, (True, True)
)

with self.assertRaisesRegex(ValueError, "expected grad_output to have sizes"):
invalid_grad_output = torch.empty(2, 1, 1, 1, 1, device=device)
torch.ops.aten.grid_sampler_3d_backward(
invalid_grad_output, input, grid, 0, 0, False, (True, True)
)

@onlyOn(["cuda", "xpu"])
@tf32_on_and_off(0.005)
def test_grid_sample_large(self, device):
Expand Down
Loading