Distributed arrays and automatic parallelization

Distributed arrays and automatic parallelization#

JAX has three styles of multi-device distributed parallelism, which can be mixed and composed. They differ in how much the compiler automatically decides versus how much is controlled explicitly in the program:

Compiler-based automatic sharding is where you program as if using a single “global view” machine, and the compiler chooses how to shard data (with some user-provided constraints via with_sharding_constraint) and how to partition computation into per-device programs with collectives.
Explicit sharding and automatic partitioning is where you still have a global view but data shardings are explicit in JAX types, inspectable using jax.typeof. The compiler still partitions the computation.
Manual per-device programming is where you have a per-device view of data and computation, and write explicit communication collectives like jax.lax.psum.

Mode	View	Explicit sharding?	Explicit collectives?
Auto	Global	❌	❌
Explicit	Global	✅	❌
Manual	Per-device	✅	✅

Before getting into details, here’s a quick example using explicit mode. First, we create a jax.Array sharded across multiple devices:

import jax
import jax.numpy as jnp
jax.config.update('jax_num_cpu_devices', 8)

jax.set_mesh(jax.make_mesh((4, 2), ('X', 'Y')))  # explicit mode by default

x = jnp.arange(8 * 4.).reshape(8, 4)
x = jax.device_put(x, jax.P('X', 'Y'))
print(jax.typeof(x))  # f32[8@X, 4@Y]

float32[8@X,4@Y]

jax.debug.visualize_array_sharding(x)

                  
  CPU 0    CPU 1  
                  
                  
  CPU 2    CPU 3  
                  
                  
  CPU 4    CPU 5  
                  
                  
  CPU 6    CPU 7

Next, we’ll apply a computation to it and observe that the result values are stored across multiple devices too:

y = jnp.sin(x).T
print(jax.typeof(y))  # f32[4@Y, 8@X]

float32[4@Y,8@X]

The jnp.sin and transpose computations were automatically parallelized across the devices on which the input values (and output values) are stored.

To understand these modes and how to switch among them, we first need to understand meshes.

A `Mesh` is a grid of devices with named axes#

Meshes made a first appearance in Data placement, as the unit of data placement: a multi-dimensional grid of devices with named axes. Here the mesh’s structure starts to matter more: because communication happens along mesh axes, the mesh shape and device order can determine communication performance. The mesh should reflect the physical connection topology among the devices.

We distinguish between concrete and abstract meshes. An abstract mesh comprises only a shape, axis names, and axis types reflecting the mode of each axis:

from __future__ import annotations
import enum

class AbstractMesh:
  axis_sizes: tuple[int, ...]
  axis_names: tuple[str, ...]
  axis_types: tuple[AxisType, ...]

class AxisType(enum.Enum):
  Auto = enum.auto()
  Explicit = enum.auto()
  Manual = enum.auto()

# A concrete mesh additionally includes physical device objects with e.g.
# precise coordinates:

import numpy as np

class Mesh:
  devices: np.ndarray[jax.Device]
  axis_names: tuple[str, ...]
  axis_types: tuple[AxisType, ...]

  @property
  def axis_sizes(self) -> tuple[int, ...]:
    return self.devices.shape

At the top level of a program (i.e. not under a jit) we can create a concrete Mesh directly using the Mesh class constructor, which lets us specify the exact device order, or using the jax.make_mesh helper, which automatically chooses a device order by taking the underlying hardware topology into account:

mesh = jax.make_mesh((4, 2), ('X', 'Y'))
print(mesh)

Mesh('X': 4, 'Y': 2, axis_types=(Explicit, Explicit))

With jax.make_mesh, all mesh axis types default to AxisType.Explicit. Beware that the Mesh constructor’s default is AxisType.Auto instead: that’s for historical reasons, and hard to migrate, even though Explicit should morally be the default. When using the Mesh constructor directly, it’s best to pass axis_types explicitly.

To avoid threading mesh throughout your program, use jax.set_mesh to set a concrete mesh globally:

jax.set_mesh(mesh)

<jax._src.sharding_impls.set_mesh at 0x7dc6f012f2b0>

You can also use with jax.set_mesh(mesh): ... as a context manager. At the top level only, the concrete mesh can be queried using jax.get_mesh() -> jax.sharding.Mesh.

Under a jit, only the abstract mesh can be queried and changed. Use jax.sharding.get_abstract_mesh() -> jax.sharding.AbstractMesh to query the current abstract mesh, and use with jax.sharding.use_abstract_mesh(m: AbstractMesh): ... to change the abstract mesh within a context. The axis sizes, axis names, and axis types can be changed, but the total size of the mesh (i.e. the product of the axis sizes) must not change.

We haven’t explained shardings yet, but here’s a toy example of changing abstract meshes inside a jax.jit:

@jax.jit
def f(x):
  abstract_mesh = jax.sharding.AbstractMesh((8,), ('A',), (jax.sharding.AxisType.Explicit,))
  with jax.sharding.use_abstract_mesh(abstract_mesh):
    y = jax.reshard(x, jax.P('A', None))
    return y * 2

z = f(x)
print(jax.typeof(z))  # f32[8@A, 4]

float32[8@A,4]

A `Sharding` describes how array values are laid out over a `Mesh`#

A jax.sharding.Sharding describes distributed memory layout. That is, it describes how an array’s entries are stored in the physical memories of different devices, i.e. how it’s sharded over devices.

At the top level, every jax.Array has an associated Sharding. Because we set a mesh above, our arrays carry a jax.NamedSharding: a concrete Mesh together with a jax.sharding.PartitionSpec (aliased to jax.P) describing the layout over it. (With no mesh set, an array’s sharding is a SingleDeviceSharding instead, as seen in Data placement.)

print(x.sharding)
jax.debug.visualize_array_sharding(x)

NamedSharding(mesh=Mesh('X': 4, 'Y': 2, axis_types=(Explicit, Explicit)), spec=P('X', 'Y'), memory_kind=device)

                  
  CPU 0    CPU 1  
                  
                  
  CPU 2    CPU 3  
                  
                  
  CPU 4    CPU 5  
                  
                  
  CPU 6    CPU 7

Here, PartitionSpec('X', 'Y') expresses that the first and second axes of the array x are sharded over the mesh axes ‘X’ and ‘Y’, respectively. We can see how that translates to physical storage using addressable_shards:

for s in x.addressable_shards:
  print(s.device, s.data, sep='\n', end='\n\n')

cpu:0
[[0. 1.]
 [4. 5.]]

cpu:1
[[2. 3.]
 [6. 7.]]

cpu:2
[[ 8.  9.]
 [12. 13.]]

cpu:3
[[10. 11.]
 [14. 15.]]

cpu:4
[[16. 17.]
 [20. 21.]]

cpu:5
[[18. 19.]
 [22. 23.]]

cpu:6
[[24. 25.]
 [28. 29.]]

cpu:7
[[26. 27.]
 [30. 31.]]

We can use jax.device_put (or jax.reshard) to produce a new array that is sharded over the same mesh of devices but with a different layout specified by a jax.P. (jax.device_put is a runtime-level API with more features than jax.reshard.) Since we have a mesh in context, via the jax.set_mesh above, we can pass jax.P instances directly to jax.device_put:

y = jax.device_put(x, jax.P('Y', 'X'))
print(y.sharding)
jax.debug.visualize_array_sharding(y)

NamedSharding(mesh=Mesh('X': 4, 'Y': 2, axis_types=(Explicit, Explicit)), spec=P('Y', 'X'), memory_kind=device)

                                    
                                    
  CPU 0    CPU 2    CPU 4    CPU 6  
                                    
                                    
                                    
                                    
                                    
  CPU 1    CPU 3    CPU 5    CPU 7

jax.reshard spells the same relayout, and works the same way at the top level or under jax.jit:

y = jax.reshard(x, jax.P('Y', 'X'))
print(y.sharding)

NamedSharding(mesh=Mesh('X': 4, 'Y': 2, axis_types=(Explicit, Explicit)), spec=P('Y', 'X'), memory_kind=device)

y = jax.device_put(x, jax.P('X', None))
print(y.sharding)
jax.debug.visualize_array_sharding(y)

NamedSharding(mesh=Mesh('X': 4, 'Y': 2, axis_types=(Explicit, Explicit)), spec=P('X', None), memory_kind=device)

Here, because the mesh axis name ‘Y’ is not mentioned in jax.P('X', None), the array is replicated over the mesh axis ‘Y’. (As a shorthand, trailing None placeholders can be omitted, so that P(‘X’, None) here means the same thing as P(‘X’). But it doesn’t hurt to be explicit!)

for s in y.addressable_shards:
  print(s.device, s.data, sep='\n', end='\n\n')

cpu:0
[[0. 1. 2. 3.]
 [4. 5. 6. 7.]]

cpu:1
[[0. 1. 2. 3.]
 [4. 5. 6. 7.]]

cpu:2
[[ 8.  9. 10. 11.]
 [12. 13. 14. 15.]]

cpu:3
[[ 8.  9. 10. 11.]
 [12. 13. 14. 15.]]

cpu:4
[[16. 17. 18. 19.]
 [20. 21. 22. 23.]]

cpu:5
[[16. 17. 18. 19.]
 [20. 21. 22. 23.]]

cpu:6
[[24. 25. 26. 27.]
 [28. 29. 30. 31.]]

cpu:7
[[24. 25. 26. 27.]
 [28. 29. 30. 31.]]

By using tuples of axis names inside a PartitionSpec, we can shard one array axis over multiple mesh axes:

y = jax.device_put(x, jax.P(('X', 'Y')))
print(y.sharding)
jax.debug.visualize_array_sharding(y)

NamedSharding(mesh=Mesh('X': 4, 'Y': 2, axis_types=(Explicit, Explicit)), spec=P(('X', 'Y'),), memory_kind=device)

So an array’s data can be replicated over a mesh axis, or one of its array axes can be sharded over that mesh axis, but there’s another possibility too: an array can be unreduced over a mesh axis:

y = jax.device_put(x, jax.P('X', None, unreduced={'Y'}))
print(y.sharding)

NamedSharding(mesh=Mesh('X': 4, 'Y': 2, axis_types=(Explicit, Explicit)), spec=P('X', None, unreduced={'Y'}, unreduced_kind=sum), memory_kind=device)

Unreduced means that the logical value equals the distributed sum of the physical shards’ values along that axis:

for s in y.addressable_shards:
  print(s.device, s.data, sep='\n', end='\n\n')

cpu:0
[[0. 1. 0. 0.]
 [4. 5. 0. 0.]]

cpu:1
[[0. 0. 2. 3.]
 [0. 0. 6. 7.]]

cpu:2
[[ 8.  9.  0.  0.]
 [12. 13.  0.  0.]]

cpu:3
[[ 0.  0. 10. 11.]
 [ 0.  0. 14. 15.]]

cpu:4
[[16. 17.  0.  0.]
 [20. 21.  0.  0.]]

cpu:5
[[ 0.  0. 18. 19.]
 [ 0.  0. 22. 23.]]

cpu:6
[[24. 25.  0.  0.]
 [28. 29.  0.  0.]]

cpu:7
[[ 0.  0. 26. 27.]
 [ 0.  0. 30. 31.]]

Unreduced is useful for delaying distributed reductions, especially in the context of autodiff. More on that in Autodiff and sharding.

Note that because every array has its own Sharding instance, and every Sharding instance has its own Mesh instance, arrays in scope can be associated with different meshes. To illustrate, we can use jax.device_put with a full jax.NamedSharding instance argument rather than using the in-context mesh:

mesh2 = jax.make_mesh((8,), ('A',))
z = jax.device_put(x, jax.NamedSharding(mesh2, jax.P('A', None)))
print(z.sharding)
print(y.sharding)

NamedSharding(mesh=Mesh('A': 8, axis_types=(Explicit,)), spec=P('A', None), memory_kind=device)
NamedSharding(mesh=Mesh('X': 4, 'Y': 2, axis_types=(Explicit, Explicit)), spec=P('X', None, unreduced={'Y'}, unreduced_kind=sum), memory_kind=device)

Now that we understand mesh shapes, axis names, and shardings at the top level, we can look at mesh axis types and how Explicit and Auto modes differ.

Explicit sharding mode makes sharding queryable at trace time#

In explicit sharding mode, shardings are always queryable via jax.typeof, even under a jax.jit:

print(jax.typeof(x).sharding)

NamedSharding(mesh=AbstractMesh('X': 4, 'Y': 2, axis_types=(Explicit, Explicit), device_kind=cpu, num_cores=None, platform=cpu), spec=P('X', 'Y'))

jax.jit(lambda x: print(jax.typeof(x).sharding))(x)

NamedSharding(mesh=AbstractMesh('X': 4, 'Y': 2, axis_types=(Explicit, Explicit), device_kind=cpu, num_cores=None, platform=cpu), spec=P('X', 'Y'))

We also call this mode “sharding in types”.

In terms of the printed representation, the type language is roughly:

 <array_type> ::= <dtype>[<size_and_sharding>, ...]
 <size_and_sharding> ::= <size> | <size>@<MeshAxisName>

Where

The MeshAxisNames in scope are those from jax.typeof(x).sharding.mesh
Each MeshAxisName must be of Explicit axis type
Each MeshAxisName can be mentioned at most once in an array type

These shardings associated with JAX types propagate through operations. For example:

x = jax.device_put(np.arange(4).reshape(4, 1), jax.P("X", None))
y = jax.device_put(np.arange(8).reshape(1, 8), jax.P(None, "Y"))

z = x + y

print(f"{jax.typeof(x)=!s}")
print(f"{jax.typeof(y)=!s}")
print(f"{jax.typeof(z)=!s}")

jax.typeof(x)=int32[4@X,1]
jax.typeof(y)=int32[1,8@Y]
jax.typeof(z)=int32[4@X,8@Y]

We can do the same type querying under a jit:

@jax.jit
def add_arrays(x, y):
  z = x + y
  print(f"{jax.typeof(x)=!s}")
  print(f"{jax.typeof(y)=!s}")
  print(f"{jax.typeof(z)=!s}")
  return z

add_arrays(x, y)

jax.typeof(x)=int32[4@X,1]
jax.typeof(y)=int32[1,8@Y]
jax.typeof(z)=int32[4@X,8@Y]

Array([[ 0,  1,  2,  3,  4,  5,  6,  7],
       [ 1,  2,  3,  4,  5,  6,  7,  8],
       [ 2,  3,  4,  5,  6,  7,  8,  9],
       [ 3,  4,  5,  6,  7,  8,  9, 10]], dtype=int32)

Given the input and output shardings, the computation itself is automatically partitioned over devices. The compiler inserts communication operations as needed. For example:

x = jax.random.normal(jax.random.key(0), (8, 4),
                      out_sharding=jax.P('X', 'Y'))
print(jax.typeof(x))

float32[8@X,4@Y]

y = x.sum(0)
print(jax.typeof(y))

float32[4@Y]

Here, when partitioning the computation, the compiler automatically inserts communication collectives to perform the reduction:

compile_txt = jax.jit(lambda x: x.sum(0)).lower(x).compile().as_text()
print('all-reduce(' in compile_txt)

True

Result shardings follow simple rules, or error and require annotation#

Each primitive operation has a sharding propagation rule to determine the sharding of the result as a function of input shardings. If there is not an obvious output sharding, an error is raised. The goal is to get important parallelism decisions in your face, rather than hide them so you might accidentally miss them. Put another way, sharding propagation rules prefer to error and require annotation rather than falling back to arbitrarily chosen defaults.

Each op is able to implement its own sharding propagation rule, but the usual pattern is:

For each output array axis, identify it with zero or more corresponding input array axes.
If all those input axes are sharded the same as each other, shard the output axis the same way; otherwise, error (and require an explicit out_sharding argument).
After all output array axes are decided that way, if an output array sharding mentions the same mesh axis more than once, error (and require an explicit out_sharding).

Here are some example rules:

nullary ops like jnp.zeros, jnp.arange: These ops create arrays from scratch so they don’t have input shardings to propagate. Their output is unsharded by default unless overridden by the out_sharding kwarg.
unary elementwise ops like sin, exp: The output is sharded the same as the input.
binary ops (+, -, * etc.): Axis shardings of “zipped” dimensions must match (or be None). “Outer product” dimensions (dimensions that appear in only one argument) are sharded as they are in the input. If the result ends up mentioning a mesh axis more than once it’s an error.

The contraction ops like jnp.dot and jnp.einsum also have some interesting cases. For example, the result of jnp.dot(x: f32[8,4@X], y:f32[4@X,16]), where the shared contracting axis is sharded the same way, could reasonably be:

f32[8,16] (doing an all-reduce)
f32[8@X,16] (a reduce-scatter on the first axis)
f32[8,16@X] (a reduce-scatter on the second axis)
f32[8,16]{U:X} (no communication)

Instead of automatically choosing one, JAX errors in this case and requires an out_sharding be provided, e.g. jnp.dot(x, y, out_sharding=jax.P('X', None)):

x = jax.device_put(jnp.arange(8 * 4.).reshape(8, 4), jax.P(None, 'X'))
y = jax.device_put(jnp.arange(4 * 16.).reshape(4, 16), jax.P('X', None))

try:
  jnp.dot(x, y)
except Exception as e:
  print("ERROR!")
  print(e)

ERROR!
Contracting dimensions are sharded and it is ambiguous how the output should be sharded. Please specify the output sharding via the `out_sharding` parameter. Got lhs_contracting_spec=('X',) and rhs_contracting_spec=('X',)

z = jnp.dot(x, y, out_sharding=jax.P('X', None))

print(jax.typeof(z))

float32[8@X,16]

But there are other jnp.dot cases that induce communication that JAX does perform automatically. For example, jnp.dot(x:f32[8,4], y:f32[4@X,16]) results in an f32[8,16], likely by doing an all-gather on y as in FSDP.

With `@auto_axes` the compiler chooses shardings within the decorated function#

If you don’t want to specify the shardings of some intermediates, and instead want the compiler to choose them automatically, you can use the @auto_axes decorator. Under this decorator, shardings aren’t queryable using jax.typeof. More specifically, auto_axes switches some or all mesh axis types to Auto, and Auto mesh axes can’t appear in array types.

Decorating a function with @auto_axes adds an out_sharding argument to the function’s signature, so the final output sharding can be set by the caller. Alternatively, decorating with @auto_axes(out_sharding=...) specifies the final output sharding at the function definition site.

For example, when our mesh axes are Explicit, operations raise an error if the sharding propagation rules lead to an inconsistent result sharding. (Adding arrays with different shardings is fine in general, as we saw above.) Here, propagation would produce a result mentioning the mesh axis X on both of its array axes:

from jax.sharding import auto_axes, explicit_axes

x = jax.device_put(np.arange(16).reshape(4, 4), jax.P("X", None))
y = jax.device_put(np.arange(16).reshape(4, 4), jax.P(None, "X"))

try:
  x + y
except Exception as e:
  print("ERROR!")
  print(e)

ERROR!
add operation with inputs: i32[4@X,4], i32[4,4@X] produces an illegally sharded result: i32[4@X,4@X]

If we just want to specify the sharding of the result and for the compiler to handle the rest, we can use auto_axes:

@auto_axes
def add2(x, y):
  print("We're in auto-sharding mode here. This is the current mesh:\n"
        f"{jax.sharding.get_abstract_mesh()}")
  return x + y

result = add2(x, y, out_sharding=jax.P("X", None))
print(f"Result type: {jax.typeof(result)}")

We're in auto-sharding mode here. This is the current mesh:
AbstractMesh('X': 4, 'Y': 2, axis_types=(Auto, Auto), device_kind=cpu, num_cores=None, platform=cpu)
Result type: int32[4@X,4]

So auto_axes lets you add an out_sharding argument to any composition of operations.

An auto_axes-decorated function can be called when the context mesh’s axis types are Explicit or Auto, but none can be in Manual. By default it switches all mesh axis types to Auto; use axes=... to switch only a subset.

Auto sharding mode decides shardings automatically during compilation#

While the auto_axes decorator is useful for temporarily switching mesh axis types from Explicit to Auto, you can also construct a Mesh with Auto axis types, e.g. at the top level:

Auto = jax.sharding.AxisType.Auto
auto_mesh = jax.make_mesh((4, 2), ('X', 'Y'), (Auto, Auto))
jax.set_mesh(auto_mesh)

x = jax.device_put(jnp.arange(8 * 4. ).reshape(8, 4 ), jax.P(None, 'X'))
y = jax.device_put(jnp.arange(4 * 16.).reshape(4, 16), jax.P('X', None))

z = jnp.dot(x, y)  # not an error!

Instead of getting an error, the compiler decided the sharding of the result!

print(z.sharding)  # works at the top-level only (i.e. outside `jit`)

NamedSharding(mesh=Mesh('X': 4, 'Y': 2, axis_types=(Auto, Auto)), spec=P(), memory_kind=device)

Whether using top-level Auto mesh axes, or using the auto_axes decorator, you can provide hints to the compiler about how intermediates should be sharded using jax.lax.with_sharding_constraint:

@jax.jit
def f(x, y):
  z = jnp.dot(x, y)
  z = jax.lax.with_sharding_constraint(z, jax.P('X', None))
  return z

z = f(x, y)
print(z.sharding)

NamedSharding(mesh=Mesh('X': 4, 'Y': 2, axis_types=(Auto, Auto)), spec=P('X',), memory_kind=device)

It’s also valid to call jax.lax.with_sharding_constraint with Explicit mode axes; for any Explicit mesh axes, it acts like an assertion that the argument’s sharding matches the specified sharding.

You can locally switch mesh axis types to Explicit using the @explicit_axes decorator:

@explicit_axes
def explicit_g(y):
  print(f'mesh inside g: {jax.sharding.get_abstract_mesh()}')
  print(f'y.sharding inside g: {jax.typeof(y) = }')
  z = y * 2
  print(f'z.sharding inside g: {jax.typeof(z) = }', end='\n\n')
  return z

@jax.jit
def f(arr1):
  print(f'mesh inside f: {jax.sharding.get_abstract_mesh()}', end='\n\n')
  x = jnp.sin(arr1)
  z = explicit_g(x, in_sharding=jax.P("X", "Y"))
  return z + 1

x = jax.device_put(np.arange(16).reshape(4, 4), jax.P("X", "Y"))
f(x)

mesh inside f: AbstractMesh('X': 4, 'Y': 2, axis_types=(Auto, Auto), device_kind=cpu, num_cores=None, platform=cpu)

mesh inside g: AbstractMesh('X': 4, 'Y': 2, axis_types=(Explicit, Explicit), device_kind=cpu, num_cores=None, platform=cpu)
y.sharding inside g: jax.typeof(y) = ShapedArray(float32[4@X,4@Y])
z.sharding inside g: jax.typeof(z) = ShapedArray(float32[4@X,4@Y])

Array([[ 1.        ,  2.682942  ,  2.818595  ,  1.28224   ],
       [-0.513605  , -0.9178486 ,  0.44116902,  2.3139732 ],
       [ 2.9787164 ,  1.824237  , -0.08804226, -0.99998045],
       [-0.07314587,  1.840334  ,  2.9812148 ,  2.3005757 ]],      dtype=float32)

It’s a kind of dual to auto_axes, where you specify in_sharding rather than out_sharding.

Concrete array shardings can mention `Auto` mesh axes#

The sharding of a concrete jax.Array can be queried via x.sharding. This can only be done at the top-level. You might expect the result to be the same as the sharding associated with the value’s type, jax.typeof(x).sharding. It might not be! The concrete array sharding, x.sharding, describes the sharding along both Explicit and Auto mesh axes. It’s the sharding that the compiler eventually chose. Whereas the type-specified sharding, jax.typeof(x).sharding, only describes the sharding along Explicit mesh axes. The Auto axes are deliberately hidden from the type because they’re the purview of the compiler. We can think of the concrete array sharding as being consistent with, but more specific than, the type-specified sharding. For example:

jax.set_mesh(jax.make_mesh((4, 2), ('X', 'Y')))  # Explicit mode

def compare_shardings(x):
  print(f"=== with mesh: {jax.sharding.get_abstract_mesh()} ===")
  print(f"Concrete value sharding: {x.sharding.spec}")
  print(f"Type-specified sharding: {jax.typeof(x).sharding.spec}\n")

my_array = jnp.sin(jax.device_put(np.arange(8), jax.P("X")))
compare_shardings(my_array)

@auto_axes
def check_in_auto_context(x):
  compare_shardings(x)
  return x

check_in_auto_context(my_array, out_sharding=jax.P("X"))

=== with mesh: AbstractMesh('X': 4, 'Y': 2, axis_types=(Explicit, Explicit), device_kind=cpu, num_cores=None, platform=cpu) ===
Concrete value sharding: P('X',)
Type-specified sharding: P('X',)

=== with mesh: AbstractMesh('X': 4, 'Y': 2, axis_types=(Auto, Auto), device_kind=cpu, num_cores=None, platform=cpu) ===
Concrete value sharding: P('X',)
Type-specified sharding: P(None,)

Array([ 0.        ,  0.84147096,  0.9092974 ,  0.14112   , -0.7568025 ,
       -0.9589243 , -0.2794155 ,  0.6569866 ], dtype=float32)

Notice that at the top level, where we’re currently in a fully Explicit mesh context, the concrete array sharding and type-specified sharding agree.

But under the auto_axes decorator we’re in a fully Auto mesh context and the two shardings disagree: the type-specified sharding is P(None) whereas the concrete array sharding is P("X") (though it could be anything! It’s up to the compiler).

Manual mode lets you write explicit collectives with a per-device view of data#

Using jax.shard_map sets mesh axis types to Manual:

mesh = jax.make_mesh((4, 2), ('X', 'Y'))
jax.set_mesh(mesh)

x = jax.device_put(jnp.arange(8 * 4. ).reshape(8, 4 ), jax.P(None, 'X'))
y = jax.device_put(jnp.arange(4 * 16.).reshape(4, 16), jax.P('X', None))

@jax.shard_map(out_specs=jax.P('X', None))
def matmul(x_shard, y_shard):
  z_summand = jnp.dot(x_shard, y_shard)
  return jax.lax.psum_scatter(z_summand, 'X', tiled=True)

z = matmul(x, y)
print(jax.typeof(z))

z_ref = jnp.dot(x, y, out_sharding=jax.P('X', None))
print(jnp.allclose(z_ref, z))

float32[8@X,16]
True

For details, see the Manual parallelism with shard_map tutorial.

Device-local layout#

A sharding answers “where do an array’s values live?” at the scale of devices: how the values are distributed across memories. Zoom into any one of those memories and the same question recurs one level down: how are that shard’s values arranged within the memory, for instance row-major or column-major? That’s the array’s device-local layout. Sharding is distributed layout, and JAX’s APIs reflect this: the same argument slots that accept shardings also accept layout-and-sharding pairs. The jax.experimental.layout package provides the types.

Layout: how an array’s values are laid out within each memory in which they reside (e.g., in the memory of a single device). A typical layout specification is a minor-to-major order listing of array dimensions.
Sharding: how an array’s values are distributed across different memory spaces, such as multiple device memories.
Format: the pairing of layout and sharding, providing a complete picture of an array’s memory placement.

Two Python types come up when controlling array layouts:

The Layout class defines the in-memory layout of an array, with key attributes:
- major_to_minor: a tuple of integers specifying the dimension ordering in memory. For example, for a 2-dimensional array, (0, 1) indicates row-major layout and (1, 0) indicates column-major.
- AUTO: a special, static sentinel object that can be used with jax.jit to request that the compiler automatically determine a good layout for a compiled function’s input or output arrays.
The Format class carries both a Layout and a Sharding, with either one taking on a default value when it is not specified. (When the layout is explicitly specified, the sharding must be as well.)

JAX API functions, such as jax.jit and jax.device_put, accept Shardings for sharding control or Formats for additional layout control. They typically do not accept Layout instances directly.

Note

The examples in this section are shown with outputs from an accelerator platform. Compiler layout choices are platform-dependent (on CPU, for instance, XLA currently chooses row-major layouts across the board), so these snippets are illustrative rather than executed in place. The jax.experimental.layout APIs are also still experimental and may change.

Specifying and reading layouts#

By passing Format objects to jax.jit in place of shardings (in the in_shardings and out_shardings arguments), you can guide the compiler’s layout decisions. Similarly you can pass Formats instead of Shardings to jax.device_put to control the layout of the resulting array.

Let’s see an example that uses both explicit and automatic layouts (as in Layout.AUTO). Imagine we have two compiled functions, init_fn and apply_fn. Say we expect init_fn to be called roughly once, but apply_fn to be called on the output of init_fn many times, so that we care much more about the performance of apply_fn. We may want to have the compiler choose a good layout for apply_fn and constrain init_fn to produce arrays of such layout. We can do this as follows:

import jax, jax.numpy as jnp
from jax.experimental.layout import Layout, Format
from jax.sharding import SingleDeviceSharding
import numpy as np

def init_fn(x, y):
  return x * 2, y * 3

def apply_fn(x, y):
  return x[0, :], y[:, 0]

Since apply_fn reads a contiguous column of its second argument y, it makes sense to lay it out in column-major order (where columns are stored contiguously). Using Layout.AUTO, we can ask the compiler to infer good input layouts and see that it indeed chooses to request the second argument in column-major layout:

shape = (4 * 128, 8 * 128)
duck = jax.ShapeDtypeStruct(shape, jnp.float32)

# Compile the `apply` function with layouts inferred automatically
apply_exe = jax.jit(
    apply_fn,
    in_shardings=Format(Layout.AUTO),
    out_shardings=Format(Layout.AUTO),
).trace(duck, duck).lower().compile()

# Read back the inferred input layout
arg_formats, kwarg_formats = apply_exe.input_formats
assert len(kwarg_formats) == 0
assert arg_formats[0].layout.major_to_minor == (0, 1)
assert arg_formats[1].layout.major_to_minor == (1, 0)

We can then compile init_fn to explicitly match this layout in its outputs:

init_exe = jax.jit(init_fn, out_shardings=arg_formats).trace(
    duck, duck).lower().compile()

assert init_exe.output_formats == arg_formats

Finally we can see how the compiled apply_fn behaves when called with differently laid out input arrays. The behavior varies with whether inputs are committed (Data placement). As the following test demonstrates, if the argument arrays are committed, then the pre-compiled apply_fn requires that they match the layout determined by the compiler above. Meanwhile it accepts uncommitted arrays of any layout (including, of course, the inferred layout). In this case, the arrays may be relaid out prior to invoking the compiled computation.

def test(x, y, msg):
  print(f'-- {msg}:')
  print('x major_to_minor =', x.format.layout.major_to_minor)
  print('y major_to_minor =', y.format.layout.major_to_minor)
  try:
    apply_exe(x, y)
    print('-> `apply` called successfully')
  except ValueError as e:
    assert 'does not match' in str(e)
    print('-> error: mismatched input layouts')
  print()

dev = jax.devices()[0]

x1 = y1 = jnp.ones(shape)
test(x1, y1, 'uncommitted with mismatched layout')

x2, y2 = init_exe(x1, y1)
test(x2, y2, 'uncommitted with matching layout')

x3 = jnp.ones(shape)
y3 = jax.device_put(np.ones(shape), Format(Layout(major_to_minor=(1, 0)),
                                           SingleDeviceSharding(dev)))
test(x3, y3, 'committed with matching layout')

x4 = jnp.ones(shape)
y4 = jax.device_put(np.ones(shape), Format(Layout(major_to_minor=(0, 1)),
                                           SingleDeviceSharding(dev)))
test(x4, y4, 'committed with mismatched layout')

-- uncommitted with mismatched layout:
x major_to_minor = (0, 1)
y major_to_minor = (0, 1)
-> `apply` called successfully

-- uncommitted with matching layout:
x major_to_minor = (0, 1)
y major_to_minor = (1, 0)
-> `apply` called successfully

-- committed with matching layout:
x major_to_minor = (0, 1)
y major_to_minor = (1, 0)
-> `apply` called successfully

-- committed with mismatched layout:
x major_to_minor = (0, 1)
y major_to_minor = (0, 1)
-> error: mismatched input layouts

Constraining intermediate layouts#

We can also enforce a specific layout on an intermediate value within a JIT-compiled function using with_layout_constraint:

from jax.experimental.layout import with_layout_constraint

@jax.jit
def f(x):
  y = x.T
  # Enforce a specific layout on `y`
  y = with_layout_constraint(y, Layout(major_to_minor=(0, 1)))
  return y * 2

This is analogous to jax.lax.with_sharding_constraint, for constraining layouts rather than shardings.

Next steps#

The next page, Manual parallelism with shard_map, is the full tutorial for manual mode: per-device programming with explicit collectives, for complete control over how computation and communication are partitioned. For how sharding interacts with automatic differentiation, including the unreduced shardings introduced above, see Autodiff and sharding in the advanced autodiff docs.