Copyright 2019 Google LLC.¶

Licensed under the Apache License, Version 2.0 (the "License");

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at

https://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.

🔪 JAX - The Sharp Bits 🔪¶

levskaya@ mattjj@

When walking about the countryside of Italy, the people will not hesitate to tell you that JAX has "una anima di pura programmazione funzionale".

JAX is a language for expressing and composing transformations of numerical programs. As such it needs to control the unwanted proliferation of side-effects in its programs so that analysis and transformation of its computations remain tractable!

This requires us to write code in a functional style with explicit descriptions of how the state of a program changes, which results in several important differences to how you might be used to programming in Numpy, Tensorflow or Pytorch.

Herein we try to cover the most frequent points of trouble that users encounter when starting out in JAX.

Installs and Imports¶

🔪 In-Place Updates¶

---

In Numpy you're used to doing this:

If we try to update a JAX device array in-place, however, we get an error! (☉_☉)

What gives?!

Allowing mutation of variables in-place makes program analysis and transformation very difficult. JAX requires a pure functional expression of a numerical program.

Instead, JAX offers the functional update functions: index_update, index_add and the index helper.

NB: Fancy Indexing is not yet supported, but will likely be added to JAX soon.

️⚠️ inside jit'd code and lax.while_loop or lax.fori_loop the size of slices can't be functions of argument values but only functions of argument shapes -- the slice start indices have no such restriction. See the below Control Flow Section for more information on this limitation.

index_update¶

If the input values of index_update aren't reused, jit-compiled code will perform these operations in-place.

index_add¶

If the input values of index_update aren't reused, jit-compiled code will perform these operations in-place.

🔪 Random Numbers¶

---

If all scientific papers whose results are in doubt because of bad rand()s were to disappear from library shelves, there would be a gap on each shelf about as big as your fist. - Numerical Recipes

RNGs and State¶

You're used to stateful pseudorandom number generators (PRNGs) from numpy and other libraries, which helpfully hide a lot of details under the hood to give you a ready fountain of pseudorandomness:

Underneath the hood, numpy uses the Mersenne Twister PRNG to power its pseudorandom functions. The PRNG has a period of $2^{19937-1}$ and at any point can be described by 624 32bit unsigned ints and a position indicating how much of this "entropy" has been used up.

This pseudorandom state vector is automagically updated behind the scenes every time a random number is needed, "consuming" 2 of the uint32s in the Mersenne twister state vector:

The problem with magic PRNG state is that it's hard to reason about how it's being used and updated across different threads, processes, and devices, and it's very easy to screw up when the details of entropy production and consumption are hidden from the end user.

The Mersenne Twister PRNG is also known to have a number of problems, it has a large 2.5Kb state size, which leads to problematic initialization issues. It fails modern BigCrush tests, and is generally slow.

JAX PRNG¶

JAX instead implements an explicit PRNG where entropy production and consumption are handled by explicitly passing and iterating PRNG state. JAX uses a modern Three-fry counter-based PRNG that's splittable. That is, its design allows us to fork the PRNG state into new PRNGs for use with parallel stochastic generation.

The random state is described by two unsigned-int32s that we call a key:

JAX's random functions produce pseudorandom numbers from the PRNG state, but do not change the state!

Reusing the same state will cause sadness and monotony, depriving the enduser of lifegiving chaos:

Instead, we split the PRNG to get usable subkeys every time we need a new pseudorandom number:

We propagate the key and make new subkeys whenever we need a new random number:

We can generate more than one subkey at a time:

🔪 Control Flow¶

---

✔ python control_flow + autodiff ✔

If you just want to apply grad to your python functions, you can use regular python control-flow constructs with no problems, as if you were using Autograd (or Pytorch or TF Eager).

python control flow + JIT

Using control flow with jit is more complicated, and by default it has more constraints.

This works:

So does this:

But this doesn't, at least by default:

What gives!?

When we jit-compile a function, we usually want to compile a version of the function that works for many different argument values, so that we can cache and reuse the compiled code. That way we don't have to re-compile on each function evaluation.

For example, if we evaluate an @jit function on the array np.array([1., 2., 3.], np.float32), we might want to compile code that we can reuse to evaluate the function on np.array([4., 5., 6.], np.float32) to save on compile time.

To get a view of your Python code that is valid for many different argument values, JAX traces it on abstract values that represent sets of possible inputs. There are multiple different levels of abstraction, and different transformations use different abstraction levels.

By default, jit traces your code on the ShapedArray abstraction level, where each abstract value represents the set of all array values with a fixed shape and dtype. For example, if we trace using the abstract value ShapedArray((3,), np.float32), we get a view of the function that can be reused for any concrete value in the corresponding set of arrays. That means we can save on compile time.

But there's a tradeoff here: if we trace a Python function on a ShapedArray((), np.float32) that isn't committed to a specific concrete value, when we hit a line like if x < 3, the expression x < 3 evaluates to an abstract ShapedArray((), np.bool_) that represents the set {True, False}. When Python attempts to coerce that to a concrete True or False, we get an error: we don't know which branch to take, and can't continue tracing! The tradeoff is that with higher levels of abstraction we gain a more general view of the Python code (and thus save on re-compilations), but we require more constraints on the Python code to complete the trace.

The good news is that you can control this tradeoff yourself. By having jit trace on more refined abstract values, you can relax the traceability constraints. For example, using the static_argnums argument to jit, we can specify to trace on concrete values of some arguments. Here's that example function again:

Here's another example, this time involving a loop:

In effect, the loop gets statically unrolled. JAX can also trace at higher levels of abstraction, like Unshaped, but that's not currently the default for any transformation

️⚠️ functions with argument-value dependent shapes

These control-flow issues also come up in a more subtle way: numerical functions we want to jit can't specialize the shapes of internal arrays on argument values (specializing on argument shapes is ok). As a trivial example, let's make a function whose output happens to depend on the input variable length.

static_argnums can be handy if length in our example rarely changes, but it would be disastrous if it changed a lot!

Lastly, if your function has global side-effects, JAX's tracer can cause weird things to happen. A common gotcha is trying to print arrays inside jit'd functions:

Structured control flow primitives¶

There are more options for control flow in JAX. Say you want to avoid re-compilations but still want to use control flow that's traceable, and that avoids un-rolling large loops. then you can use these 4 structured control flow primitives:

• lax.cond will be differentiable soon
• lax.while_loop non-differentiable*
• lax.fori_loop non-differentiable*
• lax.scan will be differentiable soon

*these can in principle be made to be forward-differentiable, but this isn't on the current roadmap.

cond¶

python equivalent:

def cond(pred, true_operand, true_fun, false_operand, false_fun):
  if pred:
    return true_fun(true_operand)
  else:
    return false_fun(false_operand)

while_loop¶

python equivalent:

def while_loop(cond_fun, body_fun, init_val):
  val = init_val
  while cond_fun(val):
    val = body_fun(val)
  return val

fori_loop¶

python equivalent:

def fori_loop(start, stop, body_fun, init_val):
  val = init_val
  for i in range(start, stop):
    val = body_fun(i, val)
  return val

Summary¶

$$ \begin{array} {r|rr} \hline \ \textrm{construct} & \textrm{jit} & \textrm{grad} \\ \hline \ \textrm{if} & ❌ & ✔ \\ \textrm{for} & ✔* & ✔\\ \textrm{while} & ✔* & ✔\\ \textrm{lax.cond} & ✔ & \textrm{soon!}\\ \textrm{lax.while_loop} & ✔ & ❌\\ \textrm{lax.fori_loop} & ✔ & ❌\\ \textrm{lax.scan} & \textrm{soon!} & \textrm{soon!}\\ \hline \end{array} $$* = argument-__value__-independent loop condition - unrolls the loop

🔪 Convolutions¶

---

JAX and XLA offer the very general N-dimensional conv_general_dilated function, but it's not very obvious how to use it. We'll give some examples of the common use-cases. There are also the convenience functions lax.conv and lax.conv_general_padding for the most common kinds of convolutions.

A survey of the family of convolutional operators, a guide to convolutional arithmetic is highly recommended reading!

Let's define a simple diagonal edge kernel:

And we'll make a simple synthetic image:

lax.conv and lax.conv_with_general_padding¶

These are the simple convenience functions for convolutions

️⚠️ The convenience lax.conv, lax.conv_with_general_padding helper function assume NCHW images and IOHW kernels.

Dimension Numbers define dimensional layout for conv_general_dilated¶

The important argument is the 3-tuple of axis layout arguments: (Input Layout, Kernel Layout, Output Layout)

• N - batch dimension
• H - spatial height
• W - spatial height
• C - channel dimension
• I - kernel input channel dimension
• O - kernel output channel dimension

⚠️ To demonstrate the flexibility of dimension numbers we choose a NHWC image and HWIO kernel convention for lax.conv_general_dilated below.

SAME padding, no stride, no dilation¶

VALID padding, no stride, no dilation¶

SAME padding, 2,2 stride, no dilation¶

VALID padding, no stride, rhs kernel dilation ~ Atrous convolution (excessive to illustrate)¶

VALID padding, no stride, lhs=input dilation ~ Transposed Convolution¶

We can use the last to, for instance, implement transposed convolutions:

1D Convolutions¶

You aren't limited to 2D convolutions, a simple 1D demo is below:

3D Convolutions¶

🔪 NaNs¶

---

Debugging NaNs¶

If you want to trace where NaNs are occurring in your functions or gradients, you can turn on the NaN-checker by:

• setting the JAX_DEBUG_NANS=True environment variable.
• adding from jax.config import config and config.update("jax_debug_nans", True) near the top of your main file
• adding from jax.config import config and config.parse_flags_with_absl() to your main file, then set the option using a command-line flag like --jax_debug_nans=True.

This will cause computations to error-out immediately on production of a NaN.

⚠️ You shouldn't have the NaN-checker on if you're not debugging, as it can introduce lots of device-host round-trips and performance regressions!

Fast math mode on CPU and disabled NaN/inf handling¶

At the moment, XLA's CPU backend defaults to enabling fast math mode, which does not preserve nan/inf semantics. (The GPU backend does not use fast math by default!) If fast math mode is enabled, the semantics of inf and nan are not preserved by XLA/LLVM, and the behavior of inf/nan values is unpredictable.

To disable fast math mode on CPU, set the environment variable:

XLA_FLAGS=--xla_cpu_enable_fast_math=false

🔪 Double (64bit) precision¶

---

At the moment, JAX by default enforces single-precision numbers to mitigate the Numpy API's tendency to aggressively promote operands to double. This is the desired behavior for many machine-learning applications, but it may catch you by surprise!

To use double-precision numbers, you need to set the jax_enable_x64 configuration variable at startup.

There are a few ways to do this:

1. You can enable 64bit mode by setting the environment variable JAX_ENABLE_X64=True.

2. You can manually set the jax_enable_x64 configuration flag at startup:

# again, this only works on startup!
from jax.config import config
config.update("jax_enable_x64", True)

3. You can parse command-line flags with absl.app.run(main)

from jax.config import config
config.config_with_absl()

4. If you want JAX to run absl parsing for you, i.e. you don't want to do absl.app.run(main), you can instead use

from jax.config import config
if __name__ == '__main__':
  # calls config.config_with_absl() *and* runs absl parsing
  config.parse_flags_with_absl()

Note that #2-#4 work for any of JAX's configuration options.

We can then confirm that x64 mode is enabled:

Caveats¶

⚠️ XLA doesn't support 64-bit convolutions on all backends!

Fin.¶

---

If something's not covered here that has caused you weeping and gnashing of teeth, please let us know and we'll extend these introductory advisos!