From dd51269856c045c182f4cf8a2a38d15369090df3 Mon Sep 17 00:00:00 2001 From: Charan Shettyhalli Guruswamy Date: Wed, 12 Aug 2026 21:11:38 +0000 Subject: [PATCH] Add Neuron dtype metadata for unfold and AA upsample TestCommon.test_dtypes treats successful dtypes absent from OpInfo as unsupported. Neuron's unfold lowering accepts integral dtypes, while its antialiased bilinear upsample lowering accepts float16 and bfloat16 in addition to the default CPU dtype set. Declare the implemented Neuron forward dtype sets in native OpInfo metadata. This removes the need for a private test override and does not change runtime lowering or backward support. Test Plan: ``` python3 -m py_compile torch/testing/_internal/common_methods_invocations.py spin quicklint -- --take RUFF spin quicklint -- --skip PYREFLY uvx --python 3.10 lintrunner@0.12.7 -a --skip PYREFLY /opt/torch-neuronx/.venv/bin/pytest -o addopts= -q -v --tb=short --timeout=800 tests/pytorch_tests_native/test/test_ops.py::TestCommonNEURON::test_dtypes__upsample_bilinear2d_aa_neuron tests/pytorch_tests_native/test/test_ops.py::TestCommonNEURON::test_dtypes_nn_functional_unfold_neuron /opt/torch-neuronx/.venv/bin/pytest -o addopts= -q -v --tb=short --timeout=1200 tests/pytorch_tests_native/test/test_ops.py::TestCommonNEURON::test_dtypes_nn_functional_unfold_neuron ``` The upsample test passed in the combined run. Unfold exceeded the local 800 second timeout, then passed alone in 805.01 seconds with the 1200 second limit. Full spin quicklint and literal lintrunner -a were also run; both reached only a pre-existing Pyrefly error in unchanged binary_cmp.py. AI-assisted: Prepared with Kiro; reviewed and approved by the author. --- torch/testing/_internal/common_methods_invocations.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py index 79a7ce813be6a..2179474034eeb 100644 --- a/torch/testing/_internal/common_methods_invocations.py +++ b/torch/testing/_internal/common_methods_invocations.py @@ -16320,6 +16320,7 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k OpInfo('nn.functional.unfold', aten_name='im2col', dtypes=floating_and_complex_types_and(torch.half, torch.bfloat16, torch.bool), + dtypesIf={'neuron': all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16)}, dtypesIfCUDA=floating_and_complex_types_and(torch.half, torch.bfloat16, torch.bool), sample_inputs_func=sample_inputs_nn_unfold, # Runs very slowly on slow gradcheck - alternatively reduce input sizes @@ -16492,6 +16493,7 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k supports_forward_ad=True, supports_fwgrad_bwgrad=True, dtypes=floating_types_and(torch.uint8), + dtypesIf={'neuron': floating_types_and(torch.uint8, torch.half, torch.bfloat16)}, dtypesIfCUDA=floating_types_and(torch.half, torch.bfloat16), gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, sample_inputs_func=partial(sample_inputs_upsample_aa, 'bilinear'),