diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 4bb392c0f02..e2148a1bafb 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -546,3 +546,21 @@ def _addmm_suite() -> WebGPUTestSuite: atol=1e-4, rtol=1e-3, ) +from executorch.backends.webgpu.test.ops.test_constant_pad_nd import ( + _randn as _pad_randn, + PadModule, +) + + +@register_op_test("constant_pad_nd") +def _constant_pad_nd_suite() -> WebGPUTestSuite: + return WebGPUTestSuite( + module_factory=lambda pad: PadModule(pad), + cases=[ + Case(name="last2", construct={"pad": [1, 2]}, inputs=(InputSpec(shape=(3, 8), gen=_pad_randn),)), + Case(name="rank3", construct={"pad": [1, 1, 2, 0]}, inputs=(InputSpec(shape=(2, 4, 8), gen=_pad_randn),)), + ], + golden_dtype="float32", + atol=1e-4, + rtol=1e-3, + ) diff --git a/backends/webgpu/test/ops/test_constant_pad_nd.py b/backends/webgpu/test/ops/test_constant_pad_nd.py new file mode 100644 index 00000000000..746cdb7b10e --- /dev/null +++ b/backends/webgpu/test/ops/test_constant_pad_nd.py @@ -0,0 +1,23 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`aten.constant_pad_nd.default` module (DaViT window padding) for op-tests.""" + +import torch + + +class PadModule(torch.nn.Module): + def __init__(self, pad): + super().__init__() + self.pad = tuple(pad) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.pad(x, self.pad, value=0.0) + + +def _randn(shape) -> torch.Tensor: + g = torch.Generator().manual_seed(sum(int(x) for x in shape)) + return torch.randn(*shape, generator=g) * 0.1