ENH: pad: normalize pad_width on the host for torch#832
Conversation
`pad_width` is host metadata (an int or (before, after) pairs), but the torch delegation branch normalized it by round-tripping through a tensor: `xp.asarray(pad_width)` places it on the *default* device, and the `tuple(...)` call to build the python ints that `torch.nn.functional.pad` expects then performs one device-to-host transfer per element. With a CUDA default device that is a needless host-to-device copy plus 2*ndim synchronizing `.item()` reads per call; with a data-free `meta` default device it raises outright. Factor the pure-python normalization out of `_funcs.pad` into `normalize_pad_width` and reuse it in the torch branch, reversing the axis order on the host. This also aligns the torch branch with the documented `pad_width` types (the tensor broadcast incidentally accepted undocumented forms such as shape ``(n, 1)``). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
| ) | ||
|
|
||
|
|
||
| def normalize_pad_width( |
There was a problem hiding this comment.
Arguably this belongs in _helpers instead
There was a problem hiding this comment.
I'll move it if you prefer that - but are you sure? It's specific to just pad and not useful as a standalone generic helper, so keeping it right above def pad: would be my choice.
There was a problem hiding this comment.
Given that it is used in 2 separate modules, I would rather move it to _helpers, yes. Let's add it to __all__ over there, too.
Would be good to at least open an issue for this |
|
It's already covered by the generic issue for non-default devices; it's one of very many (see gh-25567) so not worth a separate issue. |
Co-authored-by: Lucas Colley <lucas.colley8@gmail.com>
Sorry, I meant an issue on this repo |
|
Ah, that seems reasonable - a generic "introduce multi-device or non-default device testing for array-api-extra" issue. I think after we have it for SciPy, it might be clearer if/how it can be done in this repo. |
|
Hmm, looks like |
it is tested otherwise coverage would fail... I think the previous commit still worked due to the helper being imported into |
|
thanks Ralf! |
2 participants
pad_widthis host metadata (an int or (before, after) pairs), but the torch delegation branch normalized it by round-tripping through a tensor:xp.asarray(pad_width)places it on the default device, and thetuple(...)call to build the python ints thattorch.nn.functional.padexpects then performs one device-to-host transfer per element. With a CUDA default device that is a needless host-to-device copy plus 2*ndim synchronizing.item()reads per call; with a data-freemetadefault device it raises outright.Factor the pure-python normalization out of
_funcs.padintonormalize_pad_widthand reuse it in the torch branch, reversing the axis order on the host.This also aligns the torch branch with the documented
pad_widthtypes (the tensor broadcast incidentally accepted undocumented forms such as shape(n, 1)).Found this when working on non-default device handling issues in
scipy; SciPy tests will cover this so I won't worry about introducing multi-device testing here right now.AI Generation Disclosure: patch generated with Claude Fable 5 (had to tweak both code comments and commit message a bit, but the hard work was done by Claude here).