Add Neuron dtype metadata for unfold and AA upsample - #97
Open
aws-charansg wants to merge 1 commit into
Open
Conversation
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.
| @@ -16324,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)}, | |||
Collaborator
There was a problem hiding this comment.
You cannot add anything neuron specific here - all of this code needs to be upstreamed to PyTorch
| @@ -16496,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)}, | |||
Collaborator
There was a problem hiding this comment.
Please write this in a way that can be upstreamed and also raise a corresponding PR to PyTorch
AbigailFernandes
requested changes
Aug 13, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:
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.
Issue #, if available:
Description of changes:
By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.