functorch

I've been thinking a lot about trying to use functional dialect of Python like Coconut[0] and Hy[1] w/ JAX so I can write functional DL code. Glad to see functorch[3] as PyTorch is the library I have the most experience with. [0] http://coconut-lang.org/ [1] https://docs.hylang.org/en/alpha/ [2] https://github.com/google/jax [3] https://github.com/pytorch/functorch

The pytorch programming model is just really hard to adapt to an XLA-like compiler. Imperative python code doesn't translate to an ML graph compiler particularly well; Jax's API is functional, so it's easier to translate to the XLA API. By contrast, torch/xla uses "lazy tensors" that record the computation graph and compile when needed. The trouble is, if the compute graph changes from run to run, you end up recompiling a lot. I guess in Jax you'd just only apply `jax.jit` to the parts where the compute graph is static? I'd be curious to see examples of how this works in practice. Fwiw, there's an offshoot of pytorch that is aiming to provide this sort of API (see https://github.com/pytorch/functorch and look at eager_compilation.py). (Disclaimer: I worked on this until quite recently.)

IMO, 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) but not the second. Disclaimer: I work on PyTorch, and Functorch more specifically, although my opinions here aren't on behalf of PyTorch.