Tensor-parallel MLP in plain JAX

import jax
from jaxviz import trace_model


# Simulate eight devices so the example runs on a CPU-only machine.
jax.config.update("jax_num_cpu_devices", 8)

import jax.numpy as jnp


Explicit = jax.sharding.AxisType.Explicit
mesh = jax.make_mesh(
    (2, 4),
    ("data", "model"),
    axis_types=(Explicit, Explicit),
    devices=jax.devices("cpu"),
)


def model(x, w1, w2):
    # w1 is split along its columns, so the activations come out split on 'model'.
    hidden = jnp.tanh(jnp.dot(x, w1))
    # w2 is split along its rows, so this contracts the axis 'hidden' is split
    # on. Every device can only compute part of the sum, and the result sharding
    # has to be stated because the compiler will not guess one.
    return jnp.dot(hidden, w2, out_sharding=jax.P("data", None))


with jax.set_mesh(mesh):
    example_args = (
        jax.device_put(jnp.ones((16, 16)), jax.P("data", None)),
        jax.device_put(jnp.ones((16, 32)), jax.P(None, "model")),
        jax.device_put(jnp.ones((32, 16)), jax.P("model", None)),
    )
    trace_model(model, *example_args, view='per_device')
jaxviz