Jan Ebert

Wait! JAX’ precision is seriously behind PyTorch here! Is PyTorch just more precise or what’s going on?

While both JAX and PyTorch have single-precision (32-bit) floating point numbers as their default dtype, NumPy uses double-precision (64-bit) floats by default.

Now, when we converted the matrix from NumPy to the respective frameworks, JAX’ jnp.array created a new array from NumPy’s, thus converting the dtype to JAX’ default. This leaves us with jax_mat.dtype being jnp.float32. However, PyTorch’s torch.from_numpy adapted the dtype exactly, which is why PyTorch had double the precision to work with.

With that knowledge, let’s make the test more fair by converting torch_mat to single-precision as well:

torch_mat = torch_mat.float()
[
    ['jax', jax_spectral_radius(jax_mat).item()],
    ['torch', torch_spectral_radius(torch_mat).item()],
]

Ahh, all is well; JAX is not lacking in terms of precision after all (at least in this small example).

Let’s check out what the square root of 2 is in single-precision to see just how precise we are:

two_f32 = np.array(2, dtype=np.float32)
[['sqrt_f32', '{:.16f}'.format(np.sqrt(two_f32))]]

Let’s finally differentiate some complex numbers. You may not ever have seen this way to differentiate in PyTorch – it’s the functional way!

Just to clarify again, being able to take this gradient is a rather recent change in PyTorch: stable since PyTorch 1.9.0, June 2021. Originally, I wanted to show here that JAX is capable of differentiating something that PyTorch cannot. The competition has caught up, though!

jax_grad = jax.grad(jax_spectral_radius)(jax_mat)

torch_mat.requires_grad_(True)
torch_rho = torch_spectral_radius(torch_mat)
torch_grad = torch.autograd.grad(torch_rho, torch_mat)

decimals = 3
[
    ['function', 'grad'],
    ['jax', round_tree(jax_grad.tolist(), decimals)],
    ['torch', round_tree(torch_grad[0].tolist(), decimals)],
]

Due to JAX doing some heavy abstract syntax tree (AST) work, it includes a nice module with tree-related functions called jax.tree_util. We will use it to conjugate the tree of gradients we obtain from jax.grad (for no special reason at all).

jax_grad = jax.grad(jax_spectral_radius)(jax_complex_mat)
jax_conj_grad = jax.tree_util.tree_map(jnp.conj, jax_grad)

torch_complex_mat.requires_grad_(True)
torch_rho = torch_spectral_radius(torch_complex_mat)
torch_grad = torch.autograd.grad(torch_rho, torch_complex_mat)

decimals = 3
[
    ['type', 'gradient'],
    None,
    ['jax', round_tree(jax_grad.tolist(), decimals)],
    ['jax conj', round_tree(jax_conj_grad.tolist(), decimals)],
    ['torch', round_tree(torch_grad[0].tolist(), decimals)],
]

Oops! It seems there is a discrepancy here…

Without going too much into the math, when optimizing a function with complex inputs and real outputs, steepest-descent algorithms need the conjugate of the complex gradient in order to walk in the correct direction. As PyTorch is a deep learning framework first-and-foremost, it conjugates its gradients by default so users can go plug-and-play when fooling around with complex numbers and optimization.

• JAX chose mathematical/API consistency (also same behavior as in Autograd³).
• However, not the gradients to use for steepest-descent optimization! (Conjugate before.)
• PyTorch has the more practical default here.

def print_hello():
    print('Hello!')  # Side effect!

jax.jit(print_hello)()

Hello!

jit_print_hello = jax.jit(print_hello)
jit_print_hello()
jit_print_hello()
print('... Hello? :(')

… Hello? :(

Before we dive in here, a quick terminology heads-up: “to JIT something”, means “to just-in-time-compile something”.

Multiple interesting things happened in these short snippets:

1. We did not get any other “Hello!” after the first one in general.
2. We did not get a second “Hello!” even though we applied jax.jit to print_hello a second time.
3. We did get a “Hello!” for the first call of the JITted function.

Let’s go through these in order:

1. The print call is called a side effect in computer science. JAX does not care for these during JIT compilation, it only cares about math – or, to be more exact, whatever comes in and what comes out of the function.
2. The reason we do not get a second “Hello!” even though we apply jax.jit again (and we would expect the side effect to happen again before the function is compiled) is because JAX caches its compilation output in the background. So if we JIT the same function twice with the same arguments, the previous compilation output will be reused.

3. JAX traces what happens inside the function on its first call, building a computational graph. This means, the first call of the function executes just like a standard Python function (though even slower due to the computational graph building).

   This also means that JITting functions that result in a large computational graph (for example a Python loop that is executed very many times) can take forever to JIT only because the first tracing of it takes so long. When you encounter this issue, you can replace your loop with control flow substitutes from the jax.lax module.

Here, we’ll see that jax.jit can also be used as a decorator. However, because we need to supply another argument to jax.jit, we cannot use it as a decorator that simply. Instead, we need to combine it with the functools.partial operator. An explanation follows after the code block.

class Counter:
    count = 0

    @functools.partial(jax.jit, static_argnums=(0,))
    def inc(self):
        self.count += 1

a = Counter()
print(a.count)
a.inc()
print(a.count)
a.inc()
print(a.count)

Applying functools.partial like this results in the following (actually anonymous) function:

# Result of `functools.partial(jax.jit, static_argnums=(0,))`.
def partial_jit(*args, **kwargs):
    return jax.jit(*args, static_argnums=(0,), **kwargs)

This new partial_jit function wraps the inc method of Counter, resulting in the equivalent of the following:

class Counter:
    count = 0

    def inc(self): self.count += 1

Counter.inc = jax.jit(Counter.inc, static_argnums=(0,))

I hope that helped make sense of the code. static_argnums basically tells jax.jit to recompile the JITted function for a different argument at that place. In return, we get some freedoms (for example, we would not be able to JIT the function otherwise). We call the arguments at the positions designated by static_argnums static from now on. More on static and non-static arguments later.

a = Counter()
print(a.count)
a.inc()
print(a.count)
a.inc()
print(a.count)

Due to self being a static argument as specified via static_argnums, the function is recompiled for a new, different self⁴.

We’ll now use Python’s timeit module to benchmark a JIT-compiled version of our old friend randn (remember we implemented this in the section on RNG in JAX). We implement a simple wrapper around it in order to initialize the JIT-compiled function before we benchmark it.

You will notice some block_until_ready calls. These are due to JAX’ asynchronous execution engine. Whenever JAX executes code on a non-host device (such as a GPU), it happens asynchronously. This means that the main thread of the program continues to run ahead with a “mock” result, also called a “future”, while the actual result is computed in the background. Only when we actually query the result will we wait until it’s available.

During benchmarking, we would get these futures immediately – that’s not much use to us. So we call the block_until_ready function in order to wait until the result of the computation is actually available. You achieve the same in PyTorch using a call to torch.cuda.synchronize.

jit_randn = jax.jit(randn, static_argnums=(1, 2))

def time_str(code, number=5000):
    # Initialize cache
    exec(code)
    return timeit.timeit(code, globals=globals(), number=number)

randn_time = time_str(
    'randn(rng_key, (100, 100))[1].block_until_ready()')
jit_randn_time = time_str(
    'jit_randn(rng_key, (100, 100))[1].block_until_ready()')
[
    ['function', 'time [s]'],
    ['randn', randn_time],
    ['jit_randn', jit_randn_time],
]

A 25 % reduction! Not bad at all, especially since we shouldn’t really have much room to optimize here.

Let’s see how PyTorch does:

np_rng = np.random.default_rng()
np_randn_time = time_str('np_rng.normal(size=(100, 100))')
torch_randn_time = time_str(
    'torch.randn((100, 100)); '
    'torch.cuda.synchronize()'
)
[['np_randn', np_randn_time], ['torch_randn', torch_randn_time]]

Apparently, PyTorch is super good at generating random numbers (3 times as fast as JIT-compiled JAX!). I did not analyze this and can’t say much more about this as there can be a myriad of reasons.

#include <stdio.h>

int sum(int limit)
{
    int i, total;

    total = 0;
    for (i = 0; i < limit; ++i) {
        total += i;
    }
    return total;
}

int main(int argc, char** argnv)
{
    printf("%d\n", sum(50000));
    return 0;
}

If you’re following along, save the above to a file called sum.c.

This is the simplest sum function we can implement. Let’s see what a modern C compiler does to this code: by outputting LLVM’s lower-level representation.

Intermediate representation (IR) is a lower-level (in this case assembly-like) representation of the program. The compiler backend LLVM uses IR to achieve portability across assembly languages.

Don’t worry too much about the vim call below – we are simply filtering the LLVM IR output so it only shows the definition of the sum function. The sed call strips trailing spaces.

clang -S -emit-llvm sum.c -O1
cat sum.ll \
    | vim - +'/^define.*sum(.*{$/,/^}$/p' \
          -es --not-a-term \
    | sed 's/ *$//'

Clang is the official C frontend⁷ of the LLVM project.

Warning: assembly-like language below! Don’t worry about reading this, I’ll give the gist of it below.

define dso_local i32 @sum(i32 %0) local_unnamed_addr #0 {
  %2 = icmp sgt i32 %0, 0
  br i1 %2, label %3, label %13

3:                                                ; preds = %1
  %4 = add i32 %0, -1
  %5 = zext i32 %4 to i33
  %6 = add i32 %0, -2
  %7 = zext i32 %6 to i33
  %8 = mul i33 %5, %7
  %9 = lshr i33 %8, 1
  %10 = trunc i33 %9 to i32
  %11 = add i32 %10, %0
  %12 = add i32 %11, -1
  br label %13

13:                                               ; preds = %3, %1
  %14 = phi i32 [ 0, %1 ], [ %12, %3 ]
  ret i32 %14
}

So, without focusing too much on the details and interpreting this intuitively: please believe me that LLVM converted our sum function that used a for-loop into… the closed-form sum formula \( n \cdot (n - 1) / 2 \) (minus instead of plus due to the limit being excluded)! Isn’t that amazing?

This form of optimization is called scalar evolution and is an induction-based technique for – as you can see – quite substantial performance improvements. If you became interested in this topic, in the next section follows the link to the source code which includes references to the papers it implements.

If you really wanted to make sure the optimization happens on the machine code level, you can compile the code to an object file and disassemble it using for example the radare2 program.

LLVM scalar evolution analysis (SCEV) source code with links to papers.

Weirdly enough, sums with a step other than 1 are not optimized even though a closed-form solution exists…

To explain a bit more, an earlier version had an int sum(int limit, int step) function that allowed a varying step size. However, LLVM did not optimize this function to the closed-form solution, even though it really should be able to (from what I could see in the comments of the scalar evolution source code).

This section is about obtaining a C timing for the sum function and can safely be skipped.

In order to get some timings for C, which does not include a nice timeit module, what follows is a benchmark program for the above sum function allowing various timing methods. The idea is to mimic timeit with this. I saved this to a file called benchmark_sum.c.

#include <stdio.h>
#include <time.h>

// The CPU cycle-based timings suffer from poor resolution compared to
// wall-time measurements.
#define MY_CLOCK 0
#define MY_CLOCK_GETTIME 1
#define MY_CLOCK_GETTIME_WALL 2
#define MY_TIME 3

#define MY_CLOCK_FUN MY_CLOCK_GETTIME_WALL

int sum(int limit)
{
    int i, total;

    total = 0;
    for (i = 0; i < limit; ++i) {
        total += i;
    }
    return total;
}

int main(int args, char** argv)
{
    double duration;
    int res;
#if MY_CLOCK_FUN == MY_CLOCK
    clock_t start_time, end_time;

    start_time = clock();
#elseif MY_CLOCK_FUN == MY_CLOCK_GETTIME
    struct timespec start_time, end_time;

    clock_gettime(CLOCK_PROCESS_CPUTIME_ID, &start_time);
#elseif MY_CLOCK_FUN == MY_CLOCK_GETTIME_WALL
    struct timespec start_time, end_time;

    clock_gettime(CLOCK_MONOTONIC, &start_time);
#elseif MY_CLOCK_FUN == MY_TIME
    time_t start_time, end_time;
#endif

    res = sum(50000);

#if MY_CLOCK_FUN == MY_CLOCK
    end_time = clock();
    duration = (double) (end_time - start_time) / CLOCKS_PER_SEC;
#elseif MY_CLOCK_FUN == MY_CLOCK_GETTIME
    clock_gettime(CLOCK_PROCESS_CPUTIME_ID, &end_time);
    duration = (
        end_time.tv_sec + 1e-9 * end_time.tv_nsec
        - start_time.tv_sec - 1e-9 * start_time.tv_nsec
        );
#elseif MY_CLOCK_FUN == MY_CLOCK_GETTIME_WALL
    clock_gettime(CLOCK_MONOTONIC, &end_time);
    duration = (
        end_time.tv_sec + 1e-9 * end_time.tv_nsec
        - start_time.tv_sec - 1e-9 * start_time.tv_nsec
        );
#elseif MY_CLOCK_FUN == MY_TIME
    end_time = time();
    duration = difftime(end_time, start_time);
#endif

    // Use `res` so the `sum` call is not optimized away.
    printf("sum %d\n", res);
    printf("dur %.17g\n", duration);
    return 0;
}

Check that we get the sum formula optimization we want. I also manually checked to make sure that main does not optimize away the sum call.

clang -S -emit-llvm benchmark_sum.c -O1
cat benchmark_sum.ll \
    | vim - +'/^define.*sum(.*{$/,/^}$/p' \
          -es --not-a-term \
    | sed 's/ *$//'

define dso_local i32 @sum(i32 %0) local_unnamed_addr #0 {
  %2 = icmp sgt i32 %0, 0
  br i1 %2, label %3, label %13

3:                                                ; preds = %1
  %4 = add i32 %0, -1
  %5 = zext i32 %4 to i33
  %6 = add i32 %0, -2
  %7 = zext i32 %6 to i33
  %8 = mul i33 %5, %7
  %9 = lshr i33 %8, 1
  %10 = trunc i33 %9 to i32
  %11 = add i32 %10, %0
  %12 = add i32 %11, -1
  br label %13

13:                                               ; preds = %3, %1
  %14 = phi i32 [ 0, %1 ], [ %12, %3 ]
  ret i32 %14
}

Execute the same number of times as we do with time_str and add up the timings.

clang -o benchmark_sum.o benchmark_sum.c -O1
./benchmark_sum.o
seq 5000 \
    | xargs -n 1 ./benchmark_sum.o \
    | awk '/dur/ {total+=$2} END {print total}'

sum 1249975000
dur 9.532824124368238e-130
4.76641e-126

I copy-pasted this number at a later location so that I was able to give a live surprise.

Back to Python, let’s see whether XLA timings match Clang. Also, let’s acknowledge Python’s miserable performance.

def py_sum_up(limit):
    return sum(range(limit))

def jax_sum_up(limit):
    return jnp.sum(jnp.arange(limit))

limit = 50000

python_time = time_str('py_sum_up(limit)')
jax_time = time_str('jax_sum_up(limit).block_until_ready()')
jit_sum_up = jax.jit(jax_sum_up, static_argnums=(0,))
jax_jit_time = time_str('jit_sum_up(limit).block_until_ready()')
[
    ['function', 'time [s]', 'result'],
    ['py_sum_up', python_time, py_sum_up(limit)],
    ['jax_sum_up', jax_time, jax_sum_up(limit).item()],
    ['jit_sum_up', jax_jit_time, jit_sum_up(limit).item()],
]

Here, we manually implement the sum formula optimization. To get a better idea how fast jit_sum_up should be in the best case.

def sum_up_const(limit):
    # Since we exclude the limit, subtract one instead of adding.
    return limit * (limit - 1) // 2

[
    ['function', 'time [s]', 'result'],
    [
        'jit_sum_up',
        jax_jit_time,
        jit_sum_up(limit).item(),
    ],
    [
        'sum_up_const',
        time_str('sum_up_const(limit)'),
        sum_up_const(limit),
    ],
]

Reason? The optimization is too slow to apply and was disabled (but only on the CPU)!⁸

C with -O1 takes around 5e-126 seconds.

Just for fun, we try the same in PyTorch, using its JIT. We need to save the following code to a file torch_sum.py so PyTorch can parse the (TorchScript) source code.

import torch

__all__ = ['torch_sum_up', 'torch_jit_sum_up']

def torch_sum_up(limit):
    return torch.sum(torch.arange(limit))

torch_jit_sum_up = torch.jit.script(torch_sum_up)

from torch_sum import *

torch_limit = torch.tensor(limit)
[
    ['function', 'time [s]', 'result'],
    [
        'torch_sum_up',
        time_str(
            'torch_sum_up(torch_limit); '
            'torch.cuda.synchronize()'
        ),
        torch_sum_up(torch_limit).item(),
    ],
    [
        'torch_jit_sum_up',
        time_str(
            'torch_jit_sum_up(torch_limit); '
            'torch.cuda.synchronize()'
        ),
        sum_up_const(torch_limit).item(),
    ],
]

If you’ve been executing jax.jit code snippets, you can find XLA IR output in the generated directory (if you have set the XLA_FLAGS as in the behind-the-scenes setup code block).

What follows is some boilerplate setup code for deep learning using JAX’ built-in example machine learning libraries jax.example_libraries.stax and jax.example_libraries.optimizers. Notice how all state is explicitly handled with these. If you don’t care for this, you can skip straight to the interesting part.

I also quickly want to mention why the code is a bit larger than it could be: I like my model to be able to work with dynamic batch sizes; however, if you follow JAX’ official example, it would seem that the batch size needs to be fixed. By implementing two little extras, we are able to handle arbitrary batch sizes. First, we implement our own small flatten function/layer that flattens the whole input (unlike stax.Flatten). Second, we apply jax.vmap to our model function. The model function is always passed the model’s parameters as well, that is why we do not want to vmap over the first argument (indicated by in_axes=(None, [...])). And that’s it!

The only slight inconvenience with this setup is that we need to add an extra size-1 batch dimension in order to handle singular inputs. You’ll see this when we test our model later.

from jax.example_libraries import stax
from jax.example_libraries import optimizers

input_shape = (2, 2)

def flatten():
    def init_fun(rng_key, input_shape):
        flattened_size = jnp.prod(jnp.array(list(input_shape)))
        return (flattened_size,), ()

    def apply_fun(params, inputs, **kwargs):
        return inputs.ravel()

    return init_fun, apply_fun

def build_model(rng_key):
    model_init, model = stax.serial(
        flatten(),
        stax.Dense(64),
        stax.Relu,
        stax.Dense(64),
        stax.Relu,
        stax.Dense(64),
        stax.Relu,
        stax.Dense(1),
    )
    # Handle varying batch sizes.
    model = jax.vmap(model, in_axes=(None, 0))

    rng_key, subkey = jax.random.split(rng_key)
    output_shape, params = model_init(subkey, input_shape)

    assert output_shape == (1,)
    return rng_key, params, model

def build_opt(params):
    opt_init, update, get_params = optimizers.adam(3e-4)
    opt_state = opt_init(params)
    return opt_state, update, get_params

rng_key, params, model = build_model(rng_key)

orig_opt_state, opt_update, get_params = build_opt(params)

Here we implement our update methods. Notice how we use batched_spectral_radius instead of jit_batched_spectral_radius in order to give XLA more optimization freedom. Also, here we see conjugating the possibly complex gradients in action.

def batch_loss(params, batch):
    preds = model(params, batch)
    targets = batched_spectral_radius(batch)
    return jnp.mean(jnp.abs(preds - targets))

@jax.jit
def train_batch(step, opt_state, batch):
    params = get_params(opt_state)
    loss, grad = jax.value_and_grad(batch_loss)(params, batch)
    # Conjugate gradient for steepest-descent optimization.
    grad = jax.tree_util.tree_map(jnp.conj, grad)
    opt_state = opt_update(step, grad, opt_state)
    return opt_state, loss

A really simple deep learning training loop. We generate our batches on-demand, taking care to update our rng_key, of course!

opt_state = orig_opt_state
batch_size = 64
batch_shape = (batch_size,) + input_shape

steps = 10000
log_step_interval = 1000
start_time = time.perf_counter()
for step in range(steps):
    rng_key, batch = jit_randn(rng_key, batch_shape, dtype=complex)
    opt_state, loss = train_batch(step, opt_state, batch)
    if step % log_step_interval == 0:
        print('step ', step, '; loss ', loss, sep='')

end_time = time.perf_counter()
print('Training took', end_time - start_time, 'seconds.')

step 0; loss 1.6026156
step 1000; loss 0.39599544
step 2000; loss 0.38151193
step 3000; loss 0.4386006
step 4000; loss 0.3645811
step 5000; loss 0.38383436
step 6000; loss 0.4037715
step 7000; loss 0.3104779
step 8000; loss 0.32223767
step 9000; loss 0.40970623
Training took 7.869086120001157 seconds.

Okay, that’s some sound old MLP training. Let’s get into parallelization already.

• Caution: Experimental, undocumented and not even used, much less tested, anywhere inside JAX!

• Future versions will have a function jax.distributed.initialize, working much like PyTorch’s

  torch.distributed.init_process_group(
      [...],
      init_method="tcp://[...]",  # or "env://"
  )
  

Adapted from JAX source code (clickable version of link below):

# Adapted and compressed from
# https://github.com/google/jax/blob/4d6467727709de1c9ad220ac62783a18bcbf4990/jax/_src/distributed.py

def jax_distributed_initialize(
        coordinator_address, num_processes, process_id):
    if process_id == 0:
        global _service
        _service = \
            jax.lib.xla_extension.get_distributed_runtime_service(
                coordinator_address, num_processes)

    client = jax.lib.xla_extension.get_distributed_runtime_client(
        coordinator_address, process_id)
    client.connect()

    factory = functools.partial(
        jax.lib.xla_client.make_gpu_client, client, process_id)
    jax.lib.xla_bridge.register_backend_factory(
        'gpu', factory, priority=300)

Equivalent of the following TensorFlow:

import tensorflow as tf

gpus = tf.config.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

In JAX:

import os

# Before first JAX computation.
os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'

With all of that out of the way, let’s finally write some distributed training code!

… What’s that? We only need to add a single line in our train_batch function?

Well, almost; we also need to apply the titular jax.pmap function transformation. I’ll explain what’s going on here after the code block.

batch_axis = 'batch'

def distributed_train_batch(step, opt_state, batch):
    params = get_params(opt_state)
    loss, grad = jax.value_and_grad(batch_loss)(params, batch)
    # This is the only line we had to add to `train_batch`.
    loss, grad = jax.lax.pmean((loss, grad), batch_axis)
    # Conjugate gradient for steepest-descent optimization.
    grad = jax.tree_util.tree_map(jnp.conj, grad)
    opt_state = opt_update(step, grad, opt_state)
    return opt_state, loss

# `pmap` also `jit`s our function.
pmap_train_batch = jax.pmap(
    distributed_train_batch,
    batch_axis,
    in_axes=(None, None, 0),
    out_axes=(None, None),
)

This is already all the magic behind Horovod! (Ignoring the super-optimized communication.)

The first thing that should have caught your eye is the definition of batch_axis at the very top of the code block. This batch_axis is passed to both jax.lax.pmean and jax.pmap as the reduction axis and mapped-over axis, respectively. We need to do this because – and I hope you’ve started to notice a pattern here – jax.pmap is also nestable! By having to specify an axis for jax.pmap, the reduction operations in jax.lax will always have an axis to refer to. We use a string to name the axis here, but any hashable would do.

The call to jax.lax.pmean averages a tree of values over all processes. In our case, it averages the loss and gradient. We average the loss as well here because you usually want to log the globally averaged loss instead of the local loss to get a smoother overall picture (in Horovod, you need to enable this explicitly). The averaged gradient is then used to do the same update step on each process, so we don’t need any more synchronization afterwards.

jax.pmap has some other arguments here we haven’t seen yet, namely in_axes and out_axes. jax.vmap accepts these too, and they are very important! With them, you control the axis of each argument that the function transformation maps over. If you don’t want to map over an argument (maybe you are passing a single constant that you don’t want to copy until it has the size of the mapped-over axis), you specify in_axes at the position of the argument as None. We do this for the training step (a single integer) and model parameters (a tree of matrices). However, we do want to map over the batch somehow. We specify the very first axis here.

out_axes is similar, but for the output. If we wanted to collect the different outputs of the function transformation for each mapped-over input, we would specify the return values that we want to collect and the axis we want to collect them over at the corresponding positions in out_axes. Since we already reduce over the mapped-over values with the jax.lax.pmean call, we will not have multiple different outputs and thus use None just like for in_axes to prevent the collection.

Here, we reshape the original batch so we can split it over its first axis as desired. In the code, we obtain the world_size as the number of devices known to JAX. Going back to the multi-node setup code from before, if you initialize your distributed training like that, i.e. by starting multiple Python processes, you may want to use world_size = jax.process_count() instead. Like so often, this depends on your use case.

opt_state = orig_opt_state
# Since we are only using a single JAX process with (possibly
# emulated) multiple devices, we use `jax.device_count`.
# Usually, this would be `jax.process_count`.
world_size = jax.device_count()
assert batch_size % world_size == 0
local_batch_size = batch_size // world_size

start_time = time.perf_counter()
print('Training with', world_size, '(possibly simulated) devices.')
for step in range(steps):
    rng_key, batch = jit_randn(rng_key, batch_shape, dtype=complex)
    batch = batch.reshape(
        (world_size, local_batch_size) + batch.shape[1:])
    opt_state, loss = pmap_train_batch(step, opt_state, batch)
    if step % log_step_interval == 0:
        print('step ', step, '; loss ', loss, sep='')

end_time = time.perf_counter()
print('Distributed training took', end_time - start_time, 'seconds.')

Training with 2 (possibly simulated) devices.
step 0; loss 1.6025591
step 1000; loss 0.3969112
step 2000; loss 0.38199607
step 3000; loss 0.4381593
step 4000; loss 0.36498725
step 5000; loss 0.38379186
step 6000; loss 0.40361017
step 7000; loss 0.3104655
step 8000; loss 0.32245204
step 9000; loss 0.40979913
Distributed training took 7.309056613001303 seconds.

Even though I only used 2 simulated devices, training did speed up a bit already. Nice!

Let’s test whether our model actually learned anything useful. Maybe the logged losses don’t tell the whole story. Also, please notice that I JITted the model function for inference here; I wanted to show this off as it will really speed up your inference times. (Of course, JITting did not really make sense here since I only use the function once.)

params = get_params(opt_state)
jit_model = jax.jit(model)

ceig_mat = jnp.array([[1.0, -1.0], [1.0, 1.0]])
batched_ceig_mat = jnp.expand_dims(ceig_mat, 0)

[
    ['function', 'spectral radius sample'],
    [
        'jax_spectral_radius',
        jax_spectral_radius(ceig_mat).item(),
    ],
    [
        'model',
        jit_model(params, batched_ceig_mat)[0].item(),
    ],
]