Torch.fx.Transformer – symbolically modify PyTorch modules | Hacker News Reader

Torch.fx.Transformer – symbolically modify PyTorch modules

pytorch.org

55 pointspizza7 comments

When I saw this in the release notes a couple versions ago, one of the questions it brought up for me is, "What is the right level of specificity for meta-programming abstractions?" I.e. this lets you do transformations on graphs for modules in this ecosystem, but not for arbitrary python. Does that narrower focus give a better power-to-weight ratio, or will it close off a lot of uses?

This is just speculative discussion from my own experiences, but I appreciate that the metaprogramming constructs are restricted to operate on pure functions, vs attempting to handle things like (e.g.) Python data declarations, or syntactical transformations.

For one, it's not clear to me that you should attempt to have 1 system which handles both syntax-level expansion (macros) and function transformation (what this system, and JAX, does) -- unless you're writing a Lisp. That's not Python -- so I'm sort of glad that these systems don't try to handle the former, and instead make a nice go at the latter.

The restriction to code which is "pure" (where pure is restricted to a definition which makes sense for these array languages) is also nice from a program transformation perspective. Otherwise, authors of transformations must reason about effects, etc. JAX has started to include utilities to help with this sort of reasoning for custom (read: user-defined) primitives.

Under the quick primer for graphs, in the sample code, there is a

    self.param = torch.nn.Parameter(torch.rand(3, 4))

which is then never used? Or am I missing something?

It's used in the forward pass right under it, no?

I don't mean the one under "Overview", but the section below it

I think you’re right, it’s unused. It’s not listed in the nodes below, so I’m not sure if it’s left in to illustrate that the tracer can tell it’s unused. It’s not called out though, and it has the same shape as the first example, so maybe the author copied the module and forgot to delete it

Transformations can add new nodes, but I’m not sure they can add or remove parameters?

How different is this from JAX in practice?

Reply on news.ycombinator.com