[HN Gopher] Pytrees
       ___________________________________________________________________
        
       Pytrees
        
       Author : f_devd
       Score  : 113 points
       Date   : 2023-05-22 10:05 UTC (12 hours ago)
        
 (HTM) web link (jax.readthedocs.io)
 (TXT) w3m dump (jax.readthedocs.io)
        
       | mccoyb wrote:
       | One curious thing I discovered a few months ago: you can sort of
       | hack higher-order functions into JAX by defining "Pytree
       | closures" which introspect on normal closures, and pull out the
       | JAX tracer data from the closure environment (and put it back in,
       | when tracing is required) --- and this works! You can pass these
       | Pytree closures in and out of JIT boundaries, etc.
       | 
       | I believe JAX has a utility for this somewhere, can't quite
       | remember what this is called.
       | 
       | I typically think of JAX as quite restrictive -- but I think the
       | reality is that the only real limit on expressivity is that you
       | can't dynamically allocate inside of unbounded control flow (e.g.
       | creating new allocations inside of a while loop).
        
         | patrickkidger wrote:
         | You're thinking of `jax.closure_convert`. :)
         | 
         | (Although technically that works by tracing and extracting all
         | constants from the jaxpr, rather than introspecting the
         | function's closure cells -- it sounds like your trick is the
         | latter.)
         | 
         | When you discuss dynamic allocation, I'm guessing you're mainly
         | referring to not being able to backprop through
         | `jax.lax.while_loop`. If so, you might find
         | `equinox.internal.while_loop` interesting, which is an
         | unbounded while loop that you can backprop through! The secret
         | sauce is to use a treeverse-style checkpointing scheme.
         | 
         | https://github.com/patrick-kidger/equinox/blob/f95a8ba13fb35...
        
       | time_to_smile wrote:
       | For those curious what the big deal is here: PyTrees make it
       | wildly easier to take derivatives with respect to parameters
       | involving a complex structure. This makes it much easier to
       | organize code for non-trivial models.
       | 
       | As an example: if you want to implement logistic regression in
       | JAX, you need to optimize the weights. This is easy enough since
       | this can be modeled as a single value, a matrix of weights. If
       | you want to model a 2 layer MLP, now you have to use 2 matrices
       | of weights (at least). You could treat this as two parameters to
       | your function (which makes the derivative more complicated to
       | manage) or you could concatenate the weights and split them up,
       | etc. Annoying, but managable.
       | 
       | When you get to something like a diffusion model you now need to
       | manage parameters for a variety of different, quite complex,
       | models. It really helps if you can keep track of all these
       | parameters in whatever data structure you like, but also
       | trivially just call "grad" with regard to these and get your
       | models derivative with respect to its parameters.
       | 
       | Pytrees make this incredibly simple, and is a major quality of
       | life improvement in automatic differentiation.
        
       | iNic wrote:
       | JAX's use of pytrees is great! They implemented a lot of useful
       | utility functions, namely `tree_map`, that makes working with
       | these objects easy and intuitive. I recommend looking at their
       | neural network example library "stax".
        
       | patrickkidger wrote:
       | Shameless advert -- Equinox is a neural network library for JAX
       | based entirely around pytrees:
       | 
       | https://github.com/patrick-kidger/equinox
       | 
       | (Now on 1.1k stars so it's achieved some popularity!)
       | 
       | This makes model-building elegant (IMO), without any new
       | abstractions to learn. Quite a PyTorch-like experience overall.
        
         | Q6T46nT668w6i3m wrote:
         | You might like diffrax too. ;)
        
           | Armavica wrote:
           | diffrax is absolutely magical. I had to integrate a lot of
           | ODEs during my PhD, so I spent quite some time choosing and
           | tuning the scipy solvers for my problems, and I thought that
           | I came close to the fastest I could do in Python. Recently,
           | out of curiosity I rewrote a stiff system that I was studying
           | to solve it with diffrax, and was astonished when I saw it
           | being solved 150x faster.
        
       | albertzeyer wrote:
       | There is also the standalone library "tree" from DeepMind:
       | https://github.com/deepmind/tree
       | 
       | It provides similar functionality but is standalone and does not
       | depend on JAX, TF or anything else.
        
       ___________________________________________________________________
       (page generated 2023-05-22 23:01 UTC)