pymc-devs / pymc-devs/pytensor
Don't force `.cpu()` on all PyTorch outputs
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
This whole thing (i.e., calling out.cpu()) is suboptimal. I think we don't need it for JAX (which returns JAX arrays/ not numpy arrays), because np.asarray works with it, and I guess it doesn't work for torch tensors.
This should only be needed for updated shared variables where we have to convert to a common type as they could be used in multiple functions with distinct backends.
Perhaps we should expand a bit on the TorchLinker to perform the updates itself, and only force conversion when that's the case. This is already supported by Function.
Originally posted by @ricardoV94 in https://github.com/pymc-devs/pytensor/pull/1032#discussion_r1821221676
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in pytensor/link/pytorch/linker.py at the unconditional out.cpu() conversion, then read pytensor/compile/function/types.py around lines 1009-1017 to understand how Function handles updates. Determine how TorchLinker can handle updated shared variables itself while preserving common-type conversion only where needed. Done means PyTorch outputs are not forced through .cpu() unnecessarily, while shared-variable updates still work across distinct backends.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- backend, machine-learning
- Issue type
- Refactor
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100