Some issues with omnistaging

[cooijmanstim]

So I just updated JAX and discovered omnistaging by way of it breaking my code (which is fine). After learning more about it and figuring out what it would take to make it work for me, I decided to disable it. The pros of omnistaging, as I understand them, include trace time efficiency and code simplification. I'm desperately in favor of those, but I also find the cons hard to swallow: it breaks valid use cases with no simple workaround, and it fundamentally alters the way JAX works by introducing global mutable state.

Concretely, in my code I often use this sort of pattern to discover and initialize parameters on the fly in an initial dummy pass through the model:

import jax, jax.numpy as jp, jax.random as rax

init_rng = rax.PRNGKey(1)
params = dict()
def fn(h, x):
  global init_rng
  init_rng, rng = rax.split(init_rng)
  w = params.setdefault("w", rax.uniform(rng, shape=[]))
  return jp.tanh(h + w * x), h
jax.lax.scan(fn, 0., jp.arange(10))
print(params["w"])  # => Traced<ShapedArray(float32[])>with<DynamicJaxprTrace(level=1/0)>

With omnistaging, this doesn't work because it leaks a tracer. In fact this is pretty much one of the examples of what not to do from the omnistaging note. That note made me believe that anything that omnistage breaks was a bug to begin with, but upon reflection I don't think it's nearly so clear-cut. Omnistaging just isn't a straightforward change that only breaks already-broken code. It's a qualitative change that breaks working code by breaking the expected functional semantics. It has real implications for real things built on JAX.

How would I fix my code above? Supposedly I would use numpy instead of JAX. But then I would have to rewrite a bunch of code in numpy using numpy's random system which has an entirely different interface. Or worse, I would have to write code that calls either numpy or jax depending on context. I'm having trouble understanding why omnistaging means I can no longer just call JAX functions with constant values and get constant results.

More generally, omnistaging marks a departure from the functional style that makes JAX simple (if not always easy). Now some innocuous function like jax.numpy.zeros can mean one of two very different things depending on some mutable global state. Observe:

import jax, jax.numpy as jp

def fn(h, x):
  z = jp.zeros([])  # => Traced<ShapedArray(float32[])>with<DynamicJaxprTrace(level=1/0)>
  print(z)
  return h, x
jax.lax.scan(fn, 0., jp.arange(10))

z = jp.zeros([])
def fn(h, x):
  print(z)  # => 0.0
  return h, x
jax.lax.scan(fn, 0., jp.arange(10))

This brings back bad memories of the way Theano and Tensorflow handled symbolic control flow behind the scenes -- push some manager onto a global stack and let it take control of every operation created under its reign. This broke on-the-fly parameter construction in the same way, which is why Tensorflow has initialization infrastructure all over the place even though initialization isn't really conceptually a special operation. Theano left its users to deal with it, somehow. This is not a road anyone should wish to go down.

JAX has been a breath of fresh air with its functional style... breaking that sets an ominous precedent. Is omnistaging here to stay? Will it always be optional? Are there further departures from functional in the pipeline?

[jakevdp]

I believe outside omnistaging your code will silently do the wrong thing, because it depends on global state. Here is your original code, executed without omnistaging:

from jax import config
config.disable_omnistaging()

import jax, jax.numpy as jp, jax.random as rax

init_rng = rax.PRNGKey(1)
params = dict()

def fn(h, x):
  global init_rng
  init_rng, rng = rax.split(init_rng)
  w = params.setdefault("w", rax.uniform(rng, shape=[]))
  return jp.tanh(h + w * x), h

jax.lax.scan(fn, 0., jp.arange(10))
print(params["w"])
# 0.12568676
print(init_rng)
# [2441914641 1384938218]

And here is a snippet that performs identical logic (but that works with omnistaging):

init_rng = rax.PRNGKey(1)
params = dict()
init_rng, rng = rax.split(init_rng)
w = params.setdefault("w", rax.uniform(rng, shape=[]))

def fn(h, x):
  return jp.tanh(h + w * x), h

jax.lax.scan(fn, 0., jp.arange(10))
print(params["w"])
# 0.12568676
print(init_rng)
# [2441914641 1384938218]

The init_rng only gets split once in either case. Why? Because in traced functions, non-JAX statements only execute at trace-time, and then those values returned during the trace are used in subsequent runs of the function (See Common Gotchas: Pure Functions for more information).

A benefit of the omnistaging approach is that it tells you, loudly, that your code is not working as you intended, rather than silently giving you unintended results.

[jakevdp]

The general recommended solution for this case numpy for operations you would like to happen on the host at compile time, and use jax.numpy for operations you would like to happen on device after compilation. So, for example, if you change your dense function to this, the code snippet works regardless of whether omnistaging is enabled or disabled:

def dense(params, x, size):
  w = params.setdefault("w", np.ones([x.shape[-1], size]))
  return x @ w

Now, this doesn't immediately help your original use-case because there is no numpy counterpart to JAX's manipulation of PRNG keys; but in general this is the way we recommend approaching things.

[cooijmanstim]

Yeah I know that's the recommended workaround, but it leaves a thing or two to be desired. It means that whatever my initialization code does now needs to be in numpy, and whatever shared routines it calls now need to have duplicate implementations in numpy.

For my next trick, my real world code wraps the initialization pass in a vmap over PRNG keys to get multiple independent initializations at once (calling the non-jitted model in a Python loop takes much more time). The PRNG key I get inside the vmap is a tracer, so I couldn't use numpy even if I wanted to. Even if I could, will returning numpy arrays from vmap just work? Will it always just work (i.e. will vmap always be traced functionally)?

Now there's a limit to how complicated initialization code can get so the damage is limited too, but I don't think this is the only reasonable use case. Can you comment on why things necessarily have to be as they are? It used to be that only operations involving traced quantities got traced. Now there may or may not be some all-seeing tracer from way up the call stack that touches everything. Maybe it's a good default, but why is there no way to selectively opt out? Why can't there be a way of taking jp.ones([...]) and telling it at trace time to drop any spurious trace levels, e.g. the one introduced by scan in my example?

[jakevdp]

I don't quite understand what you mean by "wrap the initialization pass in a vmap over PRNG keys". If you're within a vmap, you'll get tracers regardless of whether omnistaging is on, no? So then how does that relate to the simpler case you started with?

To your question about why things are the way they are, I think maybe @mattjj or another dev can weigh-in. But your initial code snippet serves as a pretty good example: most Python programmers would look at that and assume the seed would be split each time the function is run; because JAX doesn't work that way, they would silently get wrong results. Omnistaging makes that failure loud, and in doing so makes JAX arguably much more user-friendly for the vast majority of users.

Does omnistaging break cases like yours, which attempt to utilize incidental side-effects to bypass the functional semantics of JAX's API? Sure. But I think use-cases that rely on side-effects are out of the scope of what JAX can support.

I would recommend redesigning your code to explicitly generate keys, and pass them as parameters to the functions that use them.

[jakevdp]

Another option we might explore is to add some sort of context manager that would allow us to explicitly mark some jax operations to not be staged, but to execute statically at trace time. That might give us the best of both worlds!

[cooijmanstim]

Sorry, I forgot about this thread. What I meant by vmapping the initialization is a structure like this:

_init_rng = None
@contextmanager
def initialization_context(rng):
  global _init_rng
  _init_rng = rng
  yield
  _init_rng = None
def pop_init_rng():
  global _init_rng
  _init_rng, rng = rax.split(_init_rng, 2)
  return rng

def initialize_matrix(m, n):
  return rax.uniform(pop_init_rng(), shape=[m, n])

def dense_layer(params, x, output_size):
  if "w" not in params:
    params["w"] = initialize_matrix(x.shape[-1], output_size)
  w = params["w"]
  return x @ w

def model(params, x):
  h = x
  for i, layer_size in enumerate([2, 3, 5]):
    h = jp.tanh(dense_layer(params.setdefault(f"layer{i}", dict()), x, layer_size))
  return h

def initialize_one(rng, x):
  with initialization_context(rng):
    params = dict()
    model(params, x)
    return params

def initialize_many(rng, x, num_replicas):
  return (jax.vmap
          (lambda rng: initialize_many(rng, x))
          (rax.split(rng, num_replicas)))

I know it doesn't look pretty, but I've been doing this kind of stuff a long time and it beats the alternatives by far. The two main objectionable mechanisms are

1. Passing initialization random key in through global contexts rather than through function arguments.
2. Passing initialized parameters out through mutation of a shared dictionary structure rather than through function return values.

Both of these can be avoided by modifying absolutely everything up the callstack to pass in an extra argument and take away an extra return value. That sounds innocuous, but it's the kind of change that turns everything it touches into lead. Every grad, scan and vmap call has to dedicate half a dozen lines to special handling of such side channels. Every lambda along the way turns into a def and reorders and dilutes the logic of the program. It's death by a thousand Pythonic cuts.

Jax is great because it is easy to reason about. Very few frameworks have this redeeming property. Theano, Tensorflow and PyTorch are all impenetrable bastions of magic, but Jax has so far escaped this fate. Jax is easy to reason about because it is functional. And so it is that I as a user can rely on its functional behavior to know that/when 1 and 2 are safe, and have half the amount of code than I otherwise would.

The opposite perspective is that the user's code should be functional so that the framework can cut corners. Should the framework serve the user or should the user serve the framework? It's a trade-off, but a rather one-sided one as any cost/benefit to the user is multiplied by the number of users. This is the whole point of code reuse.

The best part is that as far as Jax is concerned, my code is perfectly functional. I'm just fetching parameters by closure. It's none of scan's business that I create some on the fly as long as that creation doesn't computation-graphically depend on scan loop variables.

A context manager would be a very helpful escape hatch, but I'm not sure how it would work if I'm inside multiple levels of omnistaging tracing. Say a scan inside a scan, and the inner scan has some computation that depends on the outer scan's loop variables but not on the inner scan's loop variables. It seems I would have to indicate to the context manager which tracers I don't depend on.