DecompCrossRefMode not dispatching correctly

[fdrocha]

🐛 Describe the bug

I came across this while working on #85403.

I had to implement decompositions both for both upsample_bicubic2d.default and upsample_bicubic2d.vec.
The second one is CompositeExplicit and just calls the first, so it seems I should haven't had to implement it?
However, if I don't, then then default decomp is never called in the relevant tests in test/test_decomp.py.

You can check this by commenting out the upsample_bicubic2d_vec in torch/_decomp/decompositions.py in the branch above and adding an assert False to first line of upsample_bicubic2d_default. If you then run pytest tests/test_decomp.py -vk interpolate_bicubic you can see that the assert never triggers.

Any thoughts @ezyang, @ngimel, @jansel, @Chillee?

Versions

#85403

cc @Chillee @ezyang @zou3519 @albanD @samdow

[ezyang]

Unfortunately this is expected. CompositeExplicit are opaque from the perspective of torch dispatch mode, you don't get to "see through" them.

In your quoted PR, why don't you just do a straight transcription of the CompositeExplicit kernel, instead of duping the full logic?

[lezcano]

Is going from CompositeExplicit to CompositeImplicit BC breaking? I'm pretty sure it is, as every change in PyTorch, but if it's not, the best solution would be to make that change, as there's no reason for this function to be CompositeExplicit. Could we make this change @ezyang?

[fdrocha]

Unfortunately this is expected. CompositeExplicit are opaque from the perspective of torch dispatch mode, you don't get to "see through" them.

In your quoted PR, why don't you just do a straight transcription of the CompositeExplicit kernel, instead of duping the full logic?

Sorry, not sure I understand the difference between the two alternatives. Could you please clarify?

[ezyang]

It may break extension writers, but we have the latitude to make this sort of change. But I'm not sure it's right; there must have been a reason someone wrote an explicit AD formula for this case...

[lezcano]

There's no reason as far as I can tell. Both use the same backward formula. This just looks as a bad design / oversight.

@fdrocha could you make the .vec CompositeImplicit (no need to write the dispatch: part in native_functions.yaml) and remove its backward formula? And now that you're at it, can you also remove the backward overload for the .vec function.

[lezcano]

Closing this issue as it's expected behaviour.