| 65 | class TestGridPull(unittest.TestCase): |
| 66 | @parameterized.expand(TEST_1D_GP, skip_on_empty=True) |
| 67 | def test_grid_pull(self, input_param, expected): |
| 68 | result = grid_pull(**input_param) |
| 69 | if input_param["input"].requires_grad: |
| 70 | input_param["input"].retain_grad() |
| 71 | if input_param["grid"].requires_grad: |
| 72 | input_param["grid"].retain_grad() |
| 73 | if input_param["input"].requires_grad or input_param["grid"].requires_grad: |
| 74 | result.sum().backward() |
| 75 | |
| 76 | grads = [] |
| 77 | if input_param["input"].requires_grad: |
| 78 | grads.append(input_param["input"].grad.view(-1)) |
| 79 | if input_param["grid"].requires_grad: |
| 80 | grads.append(input_param["grid"].grad.view(-1)) |
| 81 | if not grads: |
| 82 | grads = torch.tensor(0.0, device=result.device) |
| 83 | elif len(grads) == 1: |
| 84 | grads = grads[0] |
| 85 | else: |
| 86 | grads = torch.cat(grads, dim=0) |
| 87 | self.assertTrue(f"{result.device}".startswith(expected["device"])) |
| 88 | np.testing.assert_allclose(result.detach().cpu().numpy(), expected["val"].cpu().numpy(), rtol=1e-4, atol=1e-4) |
| 89 | np.testing.assert_allclose(grads.detach().cpu().numpy(), expected["grad"].cpu().numpy(), rtol=1e-4, atol=1e-4) |
| 90 | |
| 91 | |
| 92 | if __name__ == "__main__": |