Jane Xu d0c103042c Introduce public torch.foreach APIs (#193607) ...

Big change! This PR is to make our foreach ops more discoverable and legit by making them public.

Design decisions:
- our new foreach ops will live in `torch.foreach`
- `torch.foreach.op` will call `torch._foreach_op`. there is no new ATen op, kernel or codegen.
- they will be behaviorally the same as our private offerings, in fact they are semantically the same in terms of overloads offered, underlying call, everything. EXCEPT we do change some param names to be more consistent and future proof
    - AND we make the API params positional/kwonly to maintain the standard for new ops.
    - AND we do not offer out variants! The existing out variants for the private offerings aren't true out variants + use more memory.
- our naming scheme: pluralize inputs (as foreach ops require some arg to be listed over) BUT maintain all other logical names (tensor1, weight, end, etc) as those can be either singular or plural.
- `torch._foreach_powsum` is intentionally private, so it is not exposed
- we wanted to keep this change minimally disruptive. we test all our new public foreach ops through the existing OpInfo mechanism + we add one-off tests to ensure that the private + public are compatible still. we do NOT run double the tests.

Invariants I reviewed:
- `_foreach_add` will route to the Tensor alpha overload when an alpha is specified and the Scalar overload when it isn't for a scalar tensor input (matching existing behavior)
- The overload orders match the order set in pythonArgParser
- default params and positionality vs keyword only match private APIs. e.g., `clone(memory_format=None)` and positional `norm(..., dtype)`
- In-place operations return the exact input list or tuple, including under torch.compile
- OpInfo names continue identifying the underlying private ATen operators  (matching aten names) so existing skips and test IDs remain valid + this change is minimal

Docs:
- I will spotcheck that the docs look good for each op -> I have now done this. I fixed up the docs a lot (will include screenshot later) where if there are multiple overloads, instead of uglily listing all of them, we prettily list them.

locally:
<img width="964" height="634" alt="image" src="https://github.com/user-attachments/assets/073514d4-d23d-4b62-a7e5-ffc07670f5fe" />

Pull Request resolved: https://github.com/pytorch/pytorch/pull/193607
Approved by: https://github.com/albanD

2026-08-21 06:13:52 +00:00

..

css

Introduce public torch.foreach APIs (#193607)

2026-08-21 06:13:52 +00:00

img

DOC: improve PyTorch logo in dark mode (#186578)

2026-06-23 19:44:58 +00:00