3 ms·
I’ve largely moved to jax, but it looks like pytorch is maybe moving that direction with torch.fx? The docs on it aren’t really clear though. Has anyone used it
by 6gvONxR4sf7o 5y ago
I’ve largely moved to jax, but it looks like pytorch is maybe moving that direction with torch.fx? The docs on it aren’t really clear though. Has anyone used it?
- chillee 5y agoIMO, FX is more of a toolkit for writing transforms over your FX modules than "moving in Jax's direction" (although there are certainly some similarities!) It's not totally clear what "Jax's direction" means to you, but I'd consider its defining characteristics as 1. composable transformations, 2. a functional way of programming (related to its function transformations) I'd say that Pytorch is moving towards the first (see https://github.com/pytorch/functorch https://github.com/pytorch/functorch) but not the second. Disclaimer: I work on PyTorch, and Functorch more specifically, although my opinions here aren't on behalf of PyTorch.
- 6gvONxR4sf7o 5y agoThat's a great answer, thank you :) Stuff like vmap and grad and pmap and all the rest have been a huge boon in simplifying some of my work, so I'm glad to see it's expanding into pytorch!