Purifying pattern

[EelcoHoogendoorn]

I was hoping to discuss what id call the 'purify pattern'. Versions of this are implemented in haiku's transform method, and also in flax.linen. The recently released ninjax provides an especially simple and general implementation of this pattern.

Ninjax can be used to write almost arbitrarily statefull and very 'torch like' code. It has some small examples but I think its more descriptive and motivating example is the dreamerv3 code.

Personally, I quite like the conceptual idea, and the elegance of the resulting code. (well most parts of it anyway)

I figured there must be some grumpy old purist functional programmers here that can put up a good rant as to why anyone would think theyd put in the time to make such a great functional programming language only for vandals to write state-like frameworks around it.

Or is that not the case, and is the purify pattern something that should actually be regarded as having a proper place in a functional language? If we called it a 'monad', would that make people feel better about it?

Jokes aside, I am curious to see some core jax developers opinions on something like the purify pattern. One thing that worries my is composability; if you look at the ninjax source it has to overload core operations like scan and grad to work properly; and with the limited unit testing in ninjax and limited experience I have using it, I am not quite confident that composing these would not come back to haunt me.

But if the purify pattern is actually a quite general pattern that is considered to be 'legit' by real functional programmers, I wonder if there isnt a case to be made to have it supported as a first class primitive to the JAX library itself?

[femtomc]

@EelcoHoogendoorn (Not a core JAX developer, but interested reader) State is not incompatible with functional programming, when done properly. Have you looked at JAX's new state module or Oryx?

To be honest with you, I'm not sure why the author of ninjax didn't create a custom interpreter (with its own Trace, Tracers, etc). The pattern there is quite bizarre to me (e.g. defining all those overloads of higher order primitives you mentioned).

I'd like to understand why the state model in this library isn't similar to Oryx's model of state - or if that model is incompatible with certain design decisions.

[EelcoHoogendoorn]

There are two design decision of ninjax that seem unique to me from the variants ive seen so far;

• a purify-scope-global rng function, so you dont need to pass rng keys around if you dont feel like it (except at the pure() boundary)
• ninjax does not have the usual init/apply split; but rather, pure() handles the initial call internally, and every invocation of the purified function explicitly passes the state pytree in and out

Not sure to what extent these features dictate the low level design decision, or to what extent they can be viewed as orthogonal high level features.

[femtomc]

They are orthogonal high level design features. Both of these features are things you could achieve with harvest (in Oryx), or an equivalent.

I’d argue harvest is more ergonomic because you don’t have to override jit, vmap, etc.

but I still do not understand the design space for ninjas — so I’m not in a place to make a judgement.

[EelcoHoogendoorn]

The JAX state module seems an internal component, not an end user api?

I was not aware of Oryx, no. Indeed it seems to also implement a form of what id call a 'purify pattern' (like flax and haiku and ninjax); lending further support to the notion that this is a wheel thatd be nice to invent once, rather than having it tied to some kind of special-purpose library.

The ninjax implementation makes quite some sense to me; but like this oryx variant, it is not exactly up to the testing and documentation standards that id expect of a first-class construct of a language.

[femtomc]

@EelcoHoogendoorn I think the intent with the state module is to eventually expose a design pattern similar to Oryx. Don't take this as truth, this is just what I've heard / inferred from various places.

[femtomc]

@EelcoHoogendoorn In general, it's hard to get this stuff right for the higher-order primitives. It requires significant expertise.

Oryx is (by a longshot) the best implementation of this concept. If you look at the code, they've covered several higher-order primitives -- and they are continually updating it to support new HOPs like pjit. Note that the people that update this in Oryx, are also the people who are working on the state module in JAX.

I think the maintainers are aware that it would be nice to have one wheel for this, as you say.

[EelcoHoogendoorn]

Nice to see the state module is under active development; though looking at the source and unit tests available at the moment I cant say that its clear to me where this is headed yet, in terms of the high level api it intends to expose.

[EelcoHoogendoorn]

@sharadmv I see you are actively working on the state module; is there a design doc pertaining to it?

[sharadmv]

Hi Eelco,

There's a WIP design doc and I will make sure to respond here once it is released (likely via JEP).

[adam-hartshorne]

How does this compare to things like Equinox?

[femtomc]

Do you mean how does ninjax compare to Equinox`? Or?

[EelcoHoogendoorn]

Yeah; seems like one could add equinox to this list of similar libraries as well; though there are a bunch of differences in the details, both in terms of the API provided as well as rng handling (flax and haiku abstract it away in their module hierarchy, ninjax provides a to-be-purified-global rng function, equinox seems to not get involved there?). I think the rng handling of ninjax is one its strong points; you only need to pass in a key at the pure() call boundary, and all code internal to that will simply tape out its 'global' rng calls within that context. By contrast I never got around to properly wrapping my head around what flax and haiku do; presumably something similar; but as part of a less generic more application-specific piece of logic I think.