Added batch rules for _upsample_bi*2d_aa and _upsample_bi*2d_aa_backward

[vfdev-5]

Description:

• Added batch rules for _upsample_bi*2d_aa and _upsample_bi*2d_aa_backward
• Added few more test cases into sample_inputs_upsample_aten

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/110172

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

⏳ No Failures, 1 Pending

As of commit 4f4dc11 with merge base 0511df0 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[kshitij12345]

I see that this is consistent with batching rule for other upsampling ops but
given that these ops are composite, I think it would make sense to decompose them instead of adding a batching rule (i.e. adding it to BatchRulesDecomposition). I am curious to know if decomposing works.

If it works, we should decompose the existing ones as well.

cc: @zou3519

Thanks!

Ref:

pytorch/aten/src/ATen/native/native_functions.yaml

Lines 12191 to 12193 in c71a64c

	- func: _upsample_bicubic2d_aa.vec(Tensor input, SymInt[]? output_size, bool align_corners, float[]? scale_factors) -> Tensor
	python_module: nn
	autogen: _upsample_bicubic2d_aa.vec_out

[vfdev-5]

Thanks Kshiteej, if I understand correctly the idea is to do

// in BatchRulesModules.cpp
EXISTING_BDIM(_upsample_bilinear2d_aa);
// and in BatchRulesDecompositions.cpp
OP_DECOMPOSE2(_upsample_bilinear2d_aa, vec);

however this leads to a compilation problem which comes from SymInt vs long types incompatibility: Native functions declares one arg as SymInt[]? but implementation uses OptionalIntArrayRef instead of OptionalSymIntArrayRef. Adapting the signature everywhere would require significant changes in this function and also underlying functions like compute_output_size which is also used in quantized upsampling... I propose to explore that and apply in a follow-up PR.

[kshitij12345]

I see, thanks for looking into it.

[zou3519]

@vfdev-5, are you planning on doing this follow-up? This PR means that a single vmap on this operator works correctly, but compositions of vmap and grad won't

[vfdev-5]

yes, let me try to do that.

[kshitij12345]

LGTM, thanks @vfdev-5

[vfdev-5]

@pytorchbot merge

[pytorchmergebot]

Merge failed

Reason: Approval needed from one of the following:
aaronenyeshi, jianyuh, bearzx, yipjustin, glaringlee, ...

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

[zou3519]

OK, but we should really decompose it, otherwise this means that it doesn't support compositions of vmap+grad

[zou3519]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here