custom_op and triton_op autograd registrations wrap user formulas in a generated autograd.Function. That wrapper passes an internal metadata object through Generated.apply and then appends a None gradient for it after calling the user backward formula.
When the user backward formula returned the wrong number of gradients, validation was left to autograd.Function after the wrapper appended the hidden metadata gradient. As a result, the error counted the internal metadata input and reported off-by-one expected/actual values.
Validate the user backward formula arity before appending the hidden metadata gradient, using the real positional operator inputs for the expected count. This keeps the internal metadata slot hidden from the diagnostic while preserving the existing metadata gradient path for valid formulas.
Fixes #137219
Generated by my agent
Test Plan:
- python test/test_custom_ops.py -k incorrect_num_gradients
- python test/test_custom_ops.py -k register_autograd
- python test/test_custom_ops.py -k supports_tensorlist
- lintrunner -a
Pull Request resolved: #186226
Approved by: https://github.com/zou3519
SOCIAL SHARE CARD GENERATOR