Autodiff and sharding

In explicit sharding mode (Distributed arrays and automatic parallelization), shardings are part of JAX types: jax.typeof(x) might print float32[8@X,4], meaning the leading axis is sharded along mesh axis X. This page is about what happens when you differentiate such a program. The main idea is that you control how your backward pass is sharded through local reasoning about your forward pass.

The rule behind this is that cotangent types are a function of primal types: the type of each gradient value in the backward pass (shape, dtype, and sharding) is determined by the type of the corresponding value in the forward pass. Once you know your forward-pass types, you know your backward-pass types, and hence where backward-pass communication happens.

There’s a second theme running through this page. If every axis of every array were plain sharded, autodiff would need nothing new: sharded cotangents for sharded primals, no communication anywhere. Replication, meaning a full copy of an array on several devices, is what requires something new. Replication’s dual under transposition is reduction, i.e. a cross-device sum, and someone has to decide where that sum happens. Demanding efficient autodiff plus local reasoning in the presence of replication is what leads to the two new sharding states introduced below, unreduced and reduced.

Most of this page works in explicit mode; the final section shows the same thing in manual mode, inside jax.shard_map().

We’ll use two CPU devices and a one-axis mesh throughout:

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

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

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

Cotangent shardings are a function of primal shardings

Here’s a data-parallel loss: the batch x is sharded along X, and the weights w are replicated (a full copy on every device):

x = jax.device_put(jnp.arange(8 * 4.).reshape(8, 4), jax.P('X', None))
w = jax.device_put(jnp.arange(4 * 2.).reshape(4, 2) / 10., jax.P(None, None))

def loss(w, x):
  return jnp.sum((x @ w) ** 2)

dw, dx = jax.grad(loss, argnums=(0, 1))(w, x)
print(jax.typeof(w), '->', jax.typeof(dw))
print(jax.typeof(x), '->', jax.typeof(dx))

float32[4,2] -> float32[4,2]
float32[8@X,4] -> float32[8@X,4]

The gradient with respect to a sharded input is sharded the same way, and the gradient with respect to a replicated input is replicated. This isn’t just true of the top-level inputs and outputs of jax.grad; it holds for every intermediate in the backward pass. We can see it in the jaxpr, where each cotangent value’s sharding matches its primal’s:

print(jax.jit(jax.grad(loss)).trace(w, x).jaxpr)

{ lambda ; a:f32[4,2] b:f32[8@X,4]. let
    c:f32[8@X,2] = dot_general[
      dimension_numbers=(([1], [0]), ([], []))
      preferred_element_type=float32
    ] b a
    d:f32[8@X,2] = integer_pow[y=2] c
    e:f32[8@X,2] = mul 2.0:f32[] c
    _:f32[] = reduce_sum[axes=(0, 1) out_sharding=None] d
    f:f32[] = reshard 1.0:f32[]
    g:f32[8@X,2] = broadcast_in_dim f
    h:f32[8@X,2] = mul g e
    i:f32[2,4] = dot_general[
      dimension_numbers=(([0], [0]), ([], []))
      preferred_element_type=float32
    ] h b
    j:f32[4,2] = transpose[permutation=(1, 0)] i
  in (j,) }

The final two equations compute dw: a dot_general of two values sharded along @X producing a replicated (unsharded) result, then a transpose. We’ll come back to what that dot costs.

There are two reasons JAX insists that cotangent shardings are determined by primal shardings:

1. User control. The goal of explicit mode is that user-written code determines all the shardings in the computation, in an easy-to-predict, local way. The backward pass is part of the computation, but it’s written by autodiff rather than by you. Making cotangent shardings a function of primal shardings means your forward-pass sharding decisions are your backward-pass sharding decisions. (Compiler-based automatic sharding mode has no analogous guarantee: there, backward-pass shardings can be chosen by the compiler, unrelated to the corresponding primal shardings.)

2. Ruling out ambiguities. If a variable is used more than once in the forward pass (fan-out), autodiff generates an addition of cotangents in the backward pass. If cotangent shardings could be unrelated to primal shardings, the two summands might have different shardings, and the addition would require communication that nothing in your code specifies. (Similarly, the zero cotangents autodiff generates would have no determined sharding.) When cotangent shardings are a function of primal shardings, both summands automatically agree: they’re cotangents for the same primal variable, so they have the same type.

This is a special case of a general principle in JAX’s autodiff: cotangent types are always a function of the corresponding primal types, meaning shapes, dtypes, and now shardings. That gives us three things: each op’s backward rule has a clear type and knows exactly what kind of cotangent it will receive; a well-typed forward program always yields a well-typed backward program; and you can predict backward-pass types by looking only at your forward-pass code, without knowing what the rest of the program looks like.

Where backward-pass communication comes from

In the example above, look at the types in the backward pass of x @ w. The cotangent dw must be replicated, because w is. But the data it’s built from is spread across devices: dw is (the transpose of) a contraction of two arrays sharded along @X, the batch axis. Producing a replicated result from a contraction over a sharded axis requires a cross-device sum, an AllReduce. The partitioner inserts it inside that final dot_general.

This is the familiar gradient synchronization of data-parallel training, and notice that you can predict it purely locally: w is replicated, each device touches only part of the batch, so somewhere in the backward pass the per-device gradient contributions must be summed. The types tell you the communication exists, but in the jaxpr above it’s implicit, hidden inside an op whose operands have a sharded contracting dimension. The rest of this page is about making that communication explicit: something you can see in the types, move around, and batch up.

Unreduced: a reduction waiting to happen

Consider a matmul whose contracting dimension is sharded:

a = jax.device_put(jnp.arange(4.).reshape(2, 2), jax.P(None, 'X'))
b = jax.device_put(jnp.arange(4., 8.).reshape(2, 2), jax.P('X', None))
print(jax.typeof(a))
print(jax.typeof(b))

float32[2,2@X]
float32[2@X,2]

a @ b

---------------------------------------------------------------------------
ShardingTypeError                         Traceback (most recent call last)
Cell In[5], line 1
----> 1 a @ b

File ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py:880, in _operator_matmul(self, other)
    878   other = m()
    879 if isinstance(other, _accepted_binop_types):
--> 880   return tensor_contractions.matmul(self, cast(ArrayLike, other))
    881 return NotImplemented

    [... skipping hidden 5 frame]

File ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/numpy/tensor_contractions.py:273, in matmul(a, b, precision, preferred_element_type, out_sharding)
    271 a = lax.squeeze(a, tuple(a_squeeze))
    272 b = lax.squeeze(b, tuple(b_squeeze))
--> 273 out = lax.dot_general(
    274   a, b, (((np.ndim(a) - 1,), (np.ndim(b) - 1 - b_is_mat,)), (a_batch, b_batch)),
    275   precision=precision, preferred_element_type=preferred_element_type,
    276   out_sharding=out_sharding)
    277 result = lax.transpose(out, perm)
    278 if both_squeeze:

    [... skipping hidden 11 frame]

File ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/lax/lax.py:5789, in _dot_general_sharding_rule(lhs, rhs, dimension_numbers, precision, preferred_element_type, out_sharding)
   5787 for l, r in zip(lhs_contracting_spec, rhs_contracting_spec):
   5788   if l is not None and r is not None:
-> 5789     raise core.ShardingTypeError(
   5790         'Contracting dimensions are sharded and it is ambiguous how the'
   5791         ' output should be sharded. Please specify the output sharding via'
   5792         ' the `out_sharding` parameter.'
   5793         f' Got {lhs_contracting_spec=} and {rhs_contracting_spec=}')
   5795 if lhs.sharding.mesh.empty and not rhs.sharding.mesh.empty:
   5796   mesh = rhs.sharding.mesh

ShardingTypeError: 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',)

JAX makes us say what we want here, because there’s a genuine choice. Look at the data each device holds:

     a: P(None, 'X')       b: P('X', None)
  (device k has column k)  (device k has row k)

        [ 0 | 1 ]             [ 4  5 ]
                      @       ---------
        [ 2 | 3 ]             [ 6  7 ]

Device 0 can multiply its column of a by its row of b without any communication, and likewise device 1. Each local matmul produces a full-shape partial sum, and the true answer is the elementwise sum of the two. One option is to finish the job with an AllReduce, by asking for an ordinary output sharding:

c = jnp.einsum('ij,jk->ik', a, b, out_sharding=jax.P(None, None))
print(jax.typeof(c))
print(c)

float32[2,2]
[[ 6.  7.]
 [26. 31.]]

The other option is to stop before the reduction:

c = jnp.einsum('ij,jk->ik', a, b,
               out_sharding=jax.P(None, None, unreduced={'X'}))
print(jax.typeof(c))

float32[2,2]{U:X}

The type float32[2,2]{U:X} reads: a 2×2 array, unreduced along mesh axis X. Each device along X holds a full-shape partial sum, and the array’s true value is the sum of those pieces:

for shard in c.addressable_shards:
  print(f'device {shard.device.id}:\n{shard.data}')

device 0:
[[ 0.  0.]
 [ 8. 10.]]
device 1:
[[ 6.  7.]
 [18. 21.]]

An unreduced array is a reduction waiting to happen. To perform it, reshard to an ordinary sharding, which runs the deferred AllReduce:

print(jax.reshard(c, jax.P(None, None)))

[[ 6.  7.]
 [26. 31.]]

Deferring lets you do more work first and pay for fewer, bigger collectives. Only linear operations make sense on unreduced arrays. Adding two arrays unreduced along the same axes is fine, since sums of partial sums are partial sums of sums:

c2 = jnp.einsum('ij,jk->ik', 2. * a, b,
                out_sharding=jax.P(None, None, unreduced={'X'}))
print(jax.typeof(c + c2))

float32[2,2]{U:X}

That lets you compute a sum of sharded matmuls (for example, a LoRA-style x @ W + x @ A @ B) with a single AllReduce at the end rather than one per matmul. Nonlinear operations, on the other hand, can’t have a rule for unreduced inputs: the cosine of a sum is not the sum of the cosines, so applying cos to each device’s partial sum would compute the wrong answer. Such ops raise an error instead:

jnp.cos(c)

---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
Cell In[11], line 1
----> 1 jnp.cos(c)

    [... skipping hidden 5 frame]

File ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/numpy/ufuncs.py:646, in cos(x)
    618 @export
    619 @jit(inline=Inline.JAX_EARLY)
    620 def cos(x: ArrayLike, /) -> Array:
    621   """Compute a trigonometric cosine of each element of input.
    622 
    623   JAX implementation of :obj:`numpy.cos`.
   (...)    644     [ 0.707 -0.    -0.707 -0.866]
    645   """
--> 646   out = lax.cos(*promote_args_inexact('cos', x))
    647   jnp_error._set_error_if_nan(out)
    648   return out

    [... skipping hidden 11 frame]

File ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/lax/lax.py:4274, in unop_ur_rule(name, aval, **kwargs)
   4272 reduced = default_unop_reduced_rule(aval)
   4273 if any(getu(aval)):
-> 4274   raise NotImplementedError(
   4275       f'unreduced rule for {name} is not implemented. Please'
   4276       ' file an issue at https://github.com/jax-ml/jax/issues')
   4277 return frozenset(), reduced, None

NotImplementedError: unreduced rule for cos is not implemented. Please file an issue at https://github.com/jax-ml/jax/issues

One caveat: for straight-line code like c + c2 above, XLA can often merge adjacent AllReduces on its own, so the compiled program may be equally good either way. The type-level guarantee matters when the compiler can’t find the merge, above all across the iterations of a loop. The microbatch example below shows that case.

Reduced: choosing unreduced gradients

Now back to autodiff. We said cotangent shardings are a function of primal shardings, and we’ve seen two entries of that function: sharded primals get sharded cotangents, replicated primals get replicated cotangents.

The replicated entry is exactly where backward-pass communication comes from: replication’s transpose is a cross-device sum. With CT(Replicated) = Replicated, autodiff performs that sum eagerly: each backward-pass op that produces a cotangent for a replicated primal does its AllReduce on the spot, implicitly, like the dw dot we saw earlier. The natural, before-any-communication state of such a cotangent is a set of per-device partial sums, an unreduced array. Suppose we want autodiff to leave it that way and let us decide when to reduce.

We need a way to ask for that in the forward pass, and it has to be a new type, since cotangent types are a function of primal types. That’s what reduced is for:

w_ = jax.reshard(w, jax.P(None, None, reduced={'X'}))
print(jax.typeof(w_))

float32[4,2]{R:X}

An array that’s reduced along X, written {R:X}, is physically identical to a replicated array: a full copy on every device along X. (The name is the past tense of “unreduced”: it’s the state an array is in after its reduction has happened.) In the forward pass it behaves exactly like a replicated array, and the reshard above moves no data. The only difference is how autodiff treats it. The complete cotangent map is:

primal type	cotangent type
sharded @X	sharded @X
replicated	replicated
reduced {R:X}	unreduced {U:X}
unreduced {U:X}	reduced {R:X}

Use Replicated and you get replicated gradients; use Reduced and you get unreduced gradients.

The reshard-to-reduced above is a communication-free cast in the forward pass, and under transposition it becomes a reshard-from-unreduced in the backward pass, which is the AllReduce. A free cast in your forward code pins down where the collective runs in your backward code, so backward-pass communication becomes something you can see and place while writing the forward pass. Compare the backward pass of our data-parallel loss with and without the cast:

def loss2(w, x):
  w = jax.reshard(w, jax.P(None, None, reduced={'X'}))
  return jnp.sum((x @ w) ** 2)

print(jax.jit(jax.grad(loss2)).trace(w, x).jaxpr)

{ lambda ; a:f32[4,2] b:f32[8@X,4]. let
    c:f32[4,2]{R:X} = reshard a
    d:f32[8@X,2] = dot_general[
      dimension_numbers=(([1], [0]), ([], []))
      preferred_element_type=float32
    ] b c
    e:f32[8@X,2] = integer_pow[y=2] d
    f:f32[8@X,2] = mul 2.0:f32[] d
    _:f32[] = reduce_sum[axes=(0, 1) out_sharding=None] e
    g:f32[] = reshard 1.0:f32[]
    h:f32[8@X,2] = broadcast_in_dim g
    i:f32[8@X,2] = mul h f
    j:f32[2,4]{U:X} = dot_general[
      dimension_numbers=(([0], [0]), ([], []))
      preferred_element_type=float32
    ] i b
    k:f32[4,2]{U:X} = transpose[permutation=(1, 0)] j
    l:f32[4,2] = reshard k
  in (l,) }

Where the original jaxpr had a dot_general with an implicit AllReduce buried inside, this one shows the dot producing an explicit f32[2,4]{U:X} value, and a final reshard (the transposed cast) turning it into the replicated dw. Same math, same total communication, but now the reduction is a visible, movable object in the program.

Once it’s visible, you can move it. Suppose the weights are used twice, as with two heads, two microbatches, or a LoRA branch:

def loss_fanout(w, x1, x2):
  w = jax.reshard(w, jax.P(None, None, reduced={'X'}))
  return jnp.sum(x1 @ w) + jnp.sum(x2 @ w)

x1 = jax.device_put(jnp.ones((8, 4)), jax.P('X', None))
x2 = jax.device_put(jnp.ones((8, 4)), jax.P('X', None))
print(jax.jit(jax.grad(loss_fanout)).trace(w, x1, x2).jaxpr)

{ lambda ; a:f32[4,2] b:f32[8@X,4] c:f32[8@X,4]. let
    d:f32[4,2]{R:X} = reshard a
    e:f32[8@X,2] = dot_general[
      dimension_numbers=(([1], [0]), ([], []))
      preferred_element_type=float32
    ] b d
    f:f32[] = reduce_sum[axes=(0, 1) out_sharding=None] e
    g:f32[8@X,2] = dot_general[
      dimension_numbers=(([1], [0]), ([], []))
      preferred_element_type=float32
    ] c d
    h:f32[] = reduce_sum[axes=(0, 1) out_sharding=None] g
    _:f32[] = add f h
    i:f32[] = reshard 1.0:f32[]
    j:f32[8@X,2] = broadcast_in_dim i
    k:f32[2,4]{U:X} = dot_general[
      dimension_numbers=(([0], [0]), ([], []))
      preferred_element_type=float32
    ] j c
    l:f32[4,2]{U:X} = transpose[permutation=(1, 0)] k
    m:f32[8@X,2] = broadcast_in_dim i
    n:f32[2,4]{U:X} = dot_general[
      dimension_numbers=(([0], [0]), ([], []))
      preferred_element_type=float32
    ] m b
    o:f32[4,2]{U:X} = transpose[permutation=(1, 0)] n
    p:f32[4,2]{U:X} = add_any l o
    q:f32[4,2] = reshard p
  in (q,) }

Both backward dots produce {U:X} contributions, the fan-out addition happens unreduced (addition is linear), and a single reshard performs one AllReduce for the whole gradient. Without the cast, each contribution would be reduced separately inside its own dot. Here the fusion is guaranteed by the types rather than left to compiler pattern-matching, which matters in the next example.

Example: microbatch gradient accumulation

The classic case is gradient accumulation. A realistic training step has two loops that the compiler can’t see through: a scan over layers inside the model, and a scan over microbatches accumulating gradients, with one update at the end. The gradient AllReduce should happen once per step, but if the weights are replicated, every microbatch’s backward pass synchronizes its own gradient contribution, inside both loops, and XLA cannot hoist collectives out of a loop for you. With reduced weights, the gradients come out unreduced, the accumulator stays unreduced across the whole scan, and you reduce once:

def predict(stacked_ws, xs):  # stacked_ws: [layer, features, features]
  def apply_layer(xs, w):
    return jnp.tanh(xs @ w), None
  final_xs, _ = jax.lax.scan(apply_layer, xs, stacked_ws)
  return final_xs

def loss3(stacked_ws, batch):
  return jnp.sum(predict(stacked_ws, batch) ** 2)

@jax.jit
def step(stacked_ws, xs):  # xs: [microbatch, batch@X, features]
  def microbatch_step(grad_acc, xs_mb):
    grads = jax.grad(loss3)(stacked_ws, xs_mb)
    # ws are reduced, so grads are unreduced -- and we can check it!
    assert jax.typeof(grads).sharding.spec.unreduced == {'X'}
    return grad_acc + grads, None

  grad_acc = jax.reshard(jnp.zeros_like(stacked_ws), jax.P(unreduced={'X'}))
  grad_acc, _ = jax.lax.scan(microbatch_step, grad_acc, xs)
  grads = jax.reshard(grad_acc, jax.P())        # the one AllReduce
  ws = jax.reshard(stacked_ws, jax.P())         # free: full copies already
  return ws - 0.01 * grads

stacked_ws = jax.device_put(jnp.stack([jnp.eye(4) / 2] * 3),
                            jax.P(reduced={'X'}))
xs = jax.device_put(jnp.ones((5, 2, 4)), jax.P(None, 'X', None))

new_ws = step(stacked_ws, xs)
print(jax.typeof(new_ws))

float32[3,4,4]

Everything type-checks locally: the weights are {R:X}, so each microbatch’s gradients come out {U:X} (even though they’re computed by a scan over layers); unreduced arrays support addition, so the microbatch scan carry accumulates them; and one reshard to replicated is the step’s single AllReduce. Notice the assert in the scan body: because shardings are part of JAX types, “the gradients are unreduced” is a property of a value that you can check with jax.typeof, inside traced code, at trace time. The updated weights come back replicated; casting them back to {R:X} for the next step is free. And the compiled program has exactly one AllReduce, outside both loops:

import re

def print_all_reduces(jitted, *args):
  hlo = jitted.lower(*args).compile().as_text()
  for line in hlo.splitlines():
    if 'all-reduce(' in line or 'all-reduce-start(' in line:
      print(re.search(r'op_name="([^"]*)"', line).group(1))

print_all_reduces(step, stacked_ws, xs)

jit(step)/reshard

Compare the same step written with plain replicated weights:

@jax.jit
def step_replicated(stacked_ws, xs):
  def microbatch_step(grad_acc, xs_mb):
    grads = jax.grad(loss3)(stacked_ws, xs_mb)  # replicated: AllReduce inside!
    return grad_acc + grads, None
  grad_acc = jnp.zeros_like(stacked_ws)
  grad_acc, _ = jax.lax.scan(microbatch_step, grad_acc, xs)
  return stacked_ws - 0.01 * grad_acc

ws_replicated = jax.reshard(stacked_ws, jax.P())
print_all_reduces(step_replicated, ws_replicated, xs)

jit(step_replicated)/while/body/eval_jaxpr/transpose(jvp())/while/body/eval_jaxpr/dot_general

The op name while/body/.../while/body shows that this AllReduce sits inside the transposed layer scan, inside the microbatch scan: it runs once per layer per microbatch. Hoisting the gradient reduction out of both loops, from once per layer per microbatch to once per step, has produced large wins in production LLM training (in one case cutting per-step time spent in gradient reduction by several times). The compiler cannot make this transformation on its own, because it can’t pattern-match collectives through a loop.

Why this design?

Some design questions that come up.

Why not make the cotangent of Replicated be Unreduced? It would be coherent, but it takes options away. Lots of code returns replicated values (loss values, for one), and making their cotangents unreduced would introduce surprising communication requirements into existing programs. By keeping Replicated ↔ Replicated and adding Reduced ↔ Unreduced as a separate pair, you get to choose per-array whether gradients arrive replicated (reduction done for you, eagerly) or unreduced (reduction is yours to place), and the choice is made by a communication-free cast in the forward pass.

Is there a theoretical justification for Unreduced, or is it just a practical trick? Here’s an argument that something like it is forced. Suppose we want the following, none of which is individually exotic:

1. autodiff should preserve communication cost, op by op: a forward op that requires no communication should transpose to a backward op that requires no communication (or at least, we need some way to write such a backward pass);

2. replication should be expressible (we don’t want only sharded axes);

3. multiplying a replicated scalar by a sharded vector, a communication-free forward op, should be expressible;

4. cotangents should have the same shape as their primals.

Transposing (3) with respect to the scalar means mapping a sharded cotangent vector to a cotangent for the replicated scalar. Doing that with no communication (1) forces a state that is scalar-shaped (4) but holds only a per-device piece of the answer, pending a sum: that’s Unreduced. (Without premise 4 you could simulate it by tacking on an extra device-indexed axis; Unreduced is in effect that axis, tracked in the sharding instead of the shape.)

Manual mode: autodiff and shard_map

In manual mode (Manual parallelism with shard_map), you write per-device code and communication is explicit: an AllReduce isn’t implied by an out_sharding or performed by a reshard; it’s a jax.lax.psum you place yourself. All the machinery above has a manual-mode counterpart, run by the same rule: cotangent types are a function of primal types.

Inside a shard_map, types track how each value relates to its counterparts on the other devices, along each mesh axis the map is manual over. There are four states per axis, mirroring the four sharding states outside (see Manual parallelism with shard_map for a fuller introduction):

outside (explicit mode)	inside (manual mode)	along X, devices hold
sharded, f32[8@X]	varying, f32[4]{V:X}	different values
replicated, f32[4]	invarying, f32[4]	the same value
unreduced, f32[4]{U:X}	unreduced, f32[4]{U:X}	partial sums of the true value
reduced, f32[4]{R:X}	reduced, f32[4]{R:X}	the same value, with unreduced cotangents

The in_specs and out_specs mediate the correspondence: a P('X') input binds to a varying value inside, P() to an invarying one, and P(unreduced={'X'}) / P(reduced={'X'}) pass through as themselves. (One difference from the outside types: sharded names the array axis that is split, as in 8@X, while inside the split has already happened, so varying is a fact about the mesh axis alone.) The cotangent map is the image of the one above: varying and invarying are each their own cotangent type, and unreduced and reduced swap.

Here’s the matmul from the unreduced section, with the same a and b, written per-device:

@jax.shard_map(in_specs=(jax.P(None, 'X'), jax.P('X', None)),
               out_specs=jax.P(None, None, unreduced={'X'}))
def matmul_partial(a, b):
  out = a @ b   # a local matmul of this device's pieces: a partial sum
  return jax.lax.pcast(out, 'X', to='unreduced')

c = matmul_partial(a, b)
print(jax.typeof(c))

float32[2,2]{U:X}

Each device multiplies its column block of a with its row block of b, a computation whose result is varying over X: each device holds a different partial sum. Declaring out_specs unreduced makes shard_map insert a free varying-to-unreduced cast, and outside we get exactly the float32[2,2]{U:X} value we built with out_sharding before. The backward pass:

def matmul_loss(a, b):
  return jnp.sum(jax.reshard(matmul_partial(a, b), jax.P(None, None)))

print(jax.jit(jax.grad(matmul_loss, argnums=(0, 1))).trace(a, b).jaxpr)

{ lambda ; a:f32[2,2@X] b:f32[2@X,2]. let
    c:f32[2,2]{U:X} = shard_map[
      check_vma=True
      in_specs=(P(None, 'X'), P('X', None))
      jaxpr={ lambda ; d:f32[2,1]{V:X} e:f32[1,2]{V:X}. let
          f:f32[2,2]{V:X} = dot_general[
            dimension_numbers=(([1], [0]), ([], []))
            preferred_element_type=float32
          ] d e
          g:f32[2,2]{U:X} = vary_unreduced_cast_p[axes=('X',)] f
        in (g,) }
      mesh=AbstractMesh('X': 2, axis_types=(Explicit,), device_kind=cpu, num_cores=None, platform=cpu)
      newly_manual_axes=frozenset({'X'})
      out_specs=(P(None, None, unreduced={'X'}, unreduced_kind=sum),)
    ] a b
    h:f32[2,2] = reshard c
    _:f32[] = reduce_sum[axes=(0, 1) out_sharding=None] h
    i:f32[] = reshard 1.0:f32[]
    j:f32[2,2] = broadcast_in_dim i
    k:f32[2,2]{R:X} = reshard j
    l:f32[2,2@X] m:f32[2@X,2] = shard_map[
      check_vma=True
      in_specs=(P(None, None, reduced={'X'}), P('X', None), P(None, 'X'))
      jaxpr={ lambda ; n:f32[2,2]{R:X} o:f32[1,2]{V:X} p:f32[2,1]{V:X}. let
          q:f32[2,2]{V:X} = reduced_vary_cast_p[axes=('X',)] n
          r:f32[2,1]{V:X} = dot_general[
            dimension_numbers=(([0], [0]), ([], []))
            preferred_element_type=float32
          ] q p
          s:f32[1,2]{V:X} = transpose[permutation=(1, 0)] r
          t:f32[2,1]{V:X} = dot_general[
            dimension_numbers=(([1], [1]), ([], []))
            preferred_element_type=float32
          ] q o
        in (t, s) }
      mesh=AbstractMesh('X': 2, axis_types=(Explicit,), device_kind=cpu, num_cores=None, platform=cpu)
      newly_manual_axes=frozenset({'X'})
      out_specs=(P(None, 'X'), P('X', None))
    ] k b a
  in (l, m) }

Read the two shard_map equations. In the forward one, the body is a local dot_general on varying operands followed by the free vary_unreduced_cast. In the backward one, the cotangent of the unreduced output arrives with in_specs=P(None, None, reduced={'X'}) (the boundary specs obey the cotangent map) and the body is reduced_vary_cast (the free transpose of the forward’s free cast) followed by two local dots. Neither body contains any communication: the single AllReduce is the transposed outer reshard, exactly where explicit mode put it.

More generally, in manual mode every type cast (jax.lax.pcast) is free and psum is the communication, and transposition pairs them up:

forward	type	transpose	type
jax.lax.psum	varying → invarying	jax.lax.pcast(..., to='varying')	invarying → varying
jax.lax.pcast(..., to='reduced')	invarying → reduced	jax.lax.psum on an unreduced value	unreduced → invarying
jax.lax.pcast(..., to='unreduced')	varying → unreduced	jax.lax.pcast(..., to='varying')	reduced → varying

The first row is the classic story (see the shard_map transposition JEP and the collectives table in Manual parallelism with shard_map): summing transposes to marking-as-varying. The second row is the manual-mode version of this page’s main trick: a free cast to reduced in the forward pass transposes to the psum that pays for it in the backward pass. And the third row is free in both directions, since declaring values to be partial sums moves no data either way. As in explicit mode, where you place free casts in your forward code determines where the backward-pass psums run.

That means the reduced-weights pattern carries over verbatim. Give the weights a reduced type on the way in, and gradients come out unreduced, with no psum anywhere in the backward pass:

w_r = jax.reshard(w, jax.P(None, None, reduced={'X'}))

@jax.shard_map(in_specs=(jax.P('X', None), jax.P(None, None, reduced={'X'})),
               out_specs=jax.P('X', None))
def apply_layer(x, w):
  return jnp.tanh(x @ w)

dw = jax.grad(lambda w, x: jnp.sum(apply_layer(x, w) ** 2))(w_r, x)
print(jax.typeof(dw))

float32[4,2]{U:X}

So a model whose layers are written with shard_map accumulates unreduced gradients across microbatches with a single AllReduce per step, exactly like the example above. For the full table of collectives and their transposes, see Manual parallelism with shard_map.