What’s going on?#

A change to JAX’s tracing infrastructure called “omnistaging”(jax-ml/jax#3370) was switched on injax==0.2.0. This change improves memory performance, trace execution time, andsimplifies jax internals, but may cause some existing code to break. Breakage isusually a result of buggy code, so long-term it’s best to fix the bugs, butomnistaging can also be disabled as a temporary workaround. And we’re happy tohelp you with fixes!

How do I know if omnistaging broke my code?#

The easiest way to tell if omnistaging is responsible is to disable omnistagingand see if the issues go away. See theWhat issues can arise when omnistagingis switched on? sectionbelow.

How can I disable omnistaging for now?#

Note: this applies to JAX versions 0.2.0 through 0.2.11; omnistaging cannot bedisabled in JAX versions 0.2.12 and higher

It is temporarily possible to disable omnistaging by

1. setting the shell environment variableJAX_OMNISTAGING to something falsey;

2. setting the boolean flagjax_omnistaging to something falsey if your codeparses flags with absl;

3. using this statement near the top of your main file:

jax.config.disable_omnistaging()

How do I fix bugs exposed by omnistaging?#

By far the most common issue with omnistaging is usingjax.numpy to computeshape values or other trace-time constants. See the code block below for a quickexample, and for full details along with other issues see the sectionWhatissues can arise when omnistaging is switchedon?.

Instead of this:

@jitdeff(x):input_size=jnp.prod(x.shape)ifinput_size>100:...

do this:

importnumpyasnp@jitdeff(x):input_size=np.prod(x.shape)ifinput_size>100:...

Instead of thinking ofjax.numpy as a drop-in replacement fornumpy, it’snow better to think of usingjax.numpy operations only when you want to perform acomputation on an accelerator (like your GPU).

Toy example#

JAX transformations likejit andpmap stage out computations to XLA. Thatis, we apply them to functions comprising multiple primitive operations so thatrather being executed one at a time from Python the operations are all part ofone end-to-end optimized XLA computation.

But exactly which operations get staged out? Until omnistaging, JAX staged outcomputation based on data dependence only. Here’s an example function, followedby the XLA HLO program it stages outbefore the omnistaging change:

fromjaximportjitimportjax.numpyasjnp@jitdeff(x):y=jnp.add(1,1)returnx*yf(3)

ENTRYjit_f.6{constant.2=pred[]constant(false)parameter.1=s32[]parameter(0)constant.3=s32[]constant(2)multiply.4=s32[]multiply(parameter.1,constant.3)ROOTtuple.5=(s32[])tuple(multiply.4)}

Notice that theadd operation is not staged out. Instead, we only see amultiply.

Here’s the HLO generated from this functionafter the omnistaging change:

ENTRYjit_f.8{constant.2=pred[]constant(false)parameter.1=s32[]parameter(0)constant.3=s32[]constant(1)constant.4=s32[]constant(1)add.5=s32[]add(constant.3,constant.4)multiply.6=s32[]multiply(parameter.1,add.5)ROOTtuple.7=(s32[])tuple(multiply.6)}

Slightly less toy example#

Here’s a less toy example which can arise in practice when we want to createboolean masks:

importjax.numpyasjnpfromjaximportlax@jitdefselect_tril(x):mask=jnp.arange(x.shape[0])[:,None]>jnp.arange(x.shape[1])returnlax.select(mask,x,jnp.zeros_like(x))# lax.select is like jnp.wherex=np.arange(12).reshape((3,4))select_tril(x)

Before omnistaging:

ENTRYjit_select_tril.8{constant.3=pred[]constant(false)constant.1=pred[3,4]{1,0}constant({...})parameter.2=s32[3,4]{1,0}parameter(0)constant.4=s32[]constant(0)broadcast.5=s32[3,4]{1,0}broadcast(constant.4),dimensions={}select.6=s32[3,4]{1,0}select(constant.1,parameter.2,broadcast.5)ROOTtuple.7=(s32[3,4]{1,0})tuple(select.6)}

Theselect operation is staged out, but the operations for constructing theconstantmask are not. Rather than being staged out, the operations thatconstructmask are executed op-by-op at Python tracing time, and XLA only seesa compile time constantconstant.1 representing the value ofmask. That’sunfortunate, because if we had staged out the operations for constructingmask, XLA could have fused them into theselect and avoided materializingthe result at all. As a result we end up wasting memory with a potentially-largeconstant, wasting time dispatching multiple un-fused op-by-op XLA computations,and potentially even fragmenting memory.

(Thebroadcast that corresponds to the construction of the zeros array forjnp.zeros_like(x) is staged out because JAX is lazy about very simpleexpressions fromjax-ml/jax#1668. Afteromnistaging, we can remove that lazy sublanguage and simplify JAX internals.)

The reason the creation ofmask is not staged out is that, before omnistaging,jit operates based on data dependence. That is,jit stages out only thoseoperations in a function that have a data dependence on an argument. Controlflow primitives andpmap behave similarly. In the case ofselect_tril, theoperations to construct the constantmask do not have a data dependence on theargument x, so they are not staged out; only thelax.select call has a datadependence.

With omnistaging alljax.numpy calls in the dynamic context of ajit-transformed function are staged out to XLA. That is, after omnistaging thecomputation XLA sees forselect_tril is

ENTRYjit_select_tril.16{constant.4=pred[]constant(false)iota.1=s32[3]{0}iota(),iota_dimension=0broadcast.5=s32[3,1]{1,0}broadcast(iota.1),dimensions={0}reshape.7=s32[3]{0}reshape(broadcast.5)broadcast.8=s32[3,4]{1,0}broadcast(reshape.7),dimensions={0}iota.2=s32[4]{0}iota(),iota_dimension=0broadcast.6=s32[1,4]{1,0}broadcast(iota.2),dimensions={1}reshape.9=s32[4]{0}reshape(broadcast.6)broadcast.10=s32[3,4]{1,0}broadcast(reshape.9),dimensions={1}compare.11=pred[3,4]{1,0}compare(broadcast.8,broadcast.10),direction=GTparameter.3=s32[3,4]{1,0}parameter(0)constant.12=s32[]constant(0)broadcast.13=s32[3,4]{1,0}broadcast(constant.12),dimensions={}select.14=s32[3,4]{1,0}select(compare.11,parameter.3,broadcast.13)ROOTtuple.15=(s32[3,4]{1,0})tuple(select.14)}

Usingjax.numpy for shape computations#

fromjaximportjitimportjax.numpyasjnp@jitdefex1(x):size=jnp.prod(jnp.array(x.shape))returnx.reshape((size,))ex1(jnp.ones((3,4)))

[... full traceback ...]  File "/home/mattjj/packages/jax/jax/core.py", line 862, in raise_concretization_error    raise ConcretizationTypeError(msg)jax.core.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected.The error arose in jax.numpy.reshape.While tracing the function ex1 at ex1.py:4, this value became a tracer due to JAX operations on these lines:  operation c:int32[] = reduce_prod[ axes=(0,) ] b:int32[2]    from line ex1.py:6 (ex1)You can use transformation parameters such as `static_argnums` for `jit` to avoid tracing particular arguments of transformed functions.See https://docs.jax.dev/en/latest/faq.html#abstract-tracer-value-encountered-where-concrete-value-is-expected-error for more information.Encountered tracer value: Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=0/1)>

Side-effects#

fromjaximportjitfromjaximportrandomkey=random.PRNGKey(0)definit():globalkeykey,subkey=random.split(key)returnrandom.normal(subkey,())print(init())# -1.2515389print(init())# -0.58665067init=jit(init)print(init())# 0.48648298print(init())# 0.48648298  !!

print(key)# Traced<ShapedArray(uint32[2])>with<DynamicJaxprTrace(level=0/1)>

random.normal(key,())

[... full stack trace …]  File "/home/mattjj/packages/jax/jax/interpreters/partial_eval.py", line 836, in _assert_live    raise core.escaped_tracer_error(msg)jax.core.UnexpectedTracerError: Encountered an unexpected tracer. Perhaps this tracer escaped through global state from a previously traced function.The functions being transformed should not save traced values to global state. Detail: tracer created on line example.py:8 (init).

Small numerical differences based on XLA optimizations#

Because with omnistaging more computations are being staged out to XLA, ratherthan some being executed at trace time, that can have the effect of reorderingfloating point operations. As a result, we’ve seen numerical behaviors change ina way that causes tests with overly tight tolerances to fail when omnistaging isswitched on.

Dependence on JAX internal APIs that changed#

Omnistaging involved some big revisions to JAX’s core code, including removingor changing internal functions. Any code that relies on such internalJAX APIs can break when omnistaging is switched on, either with build errors(from pytype) or runtime errors.

Triggering XLA compile time bugs#

Because omnistaging involves staging out more code to XLA, we’ve seen it triggerpre-existing XLA compile-time bugs on some backends. The best thing to do withthese is to report them so we can work with the XLA teams on fixes.