Rank-reducing vs rank-preserving maps over array axes#

We can think ofpmap (andvmap andxmap) as unstacking each array inputalong an axis (e.g. unpacking a 2D matrix into its 1D rows), applying its bodyfunction to each piece, and stacking the results back together, at least whencollectives aren’t involved:

pmap(f,in_axes=[0],out_axes=0)(xs)==jnp.stack([f(x)forxinxs])

For example, ifxs had shapef32[8,5] then eachx has shapef32[5], andif eachf(x) has shapef32[3,7] then the final stacked resultpmap(f)(xs)has shapef32[8,3,7]. That is, each application of the body functionf takesas argument inputs with one fewer axis than the corresponding argument topmap(f). We can say these arerank-reducing maps with unstacking/stacking ofinputs/outputs.

The number of logical applications off is determined by the size of the inputaxis being mapped over: for example, if we map over an input axis of size 8,semantically we get 8 logical applications of the function, which for pmapalways correspond to 8 devices physically computing them.

In contrast,shmap does not have this rank-reducing behavior. Instead, we canthink of it as slicing (or “unconcatenating”) along input axes into blocks,applying the body function, and concatenating the results back together (againwhen collectives aren’t involved):

devices=np.array(jax.devices()[:4])m=Mesh(devices,('i',))# mesh.shape['i'] = 4shard_map(f,m,in_specs=P('i'),out_specs=P('i'))(y)==jnp.concatenate([f(y_blk)fory_blkinjnp.split(y,4)])

Recall thatjnp.split slices its input into equally-sized blocks with the samerank, so that if in the above exampley has shapef32[8,5] then eachy_blkhas shapef32[2,5], and if eachf(y_blk) has shapef32[3,7] then the finalconcatenated resultshard_map(f,...)(y) has shapef32[12,7]. Soshmap(shard_map) maps over shards, or blocks, of its inputs. We can say it’s arank-preserving map with unconcatenating/concatenating of its inputs/outputs.

The number of logical applications off is determined by the mesh size, not byany input axis size: for example, if we have a mesh of total size 4 (i.e. over 4devices) then semantically we get 4 logical applications of the function,corresponding to the 4 devices physically computing them.

Controlling how each input is split (unconcatenated) and tiled within_specs#

Each of thein_specs identifies some of the corresponding input array’s axeswith mesh axes by name usingPartitionSpecs, representing how to split (orunconcatenate) that input into the blocks to which the body function is applied.That identification determines the shard sizes; when an input axis is identifiedwith a mesh axis, the input is split (unconcatenated) along that logical axisinto a number of pieces equal to the corresponding mesh axis size. (It’s anerror if the corresponding mesh axis size does not evenly divide the input arrayaxis size.) If an input’s pspec does not mention a mesh axis name, then there’sno splitting over that mesh axis. For example:

devices=np.array(jax.devices())m=Mesh(devices.reshape(4,2),('i','j'))@partial(shard_map,mesh=m,in_specs=P('i',None),out_specs=P('i','j'))deff1(x_block):print(x_block.shape)returnx_blockx1=np.arange(12*12).reshape(12,12)y=f1(x1)# prints (3,12)

Here, because the input pspec did not mention the mesh axis name'j', no inputarray axis is split over that mesh axis; similarly, because the second axis ofthe input array is not identified with (and hence split over) any mesh axis,application off1 gets a full view of the input along that axis.

When a mesh axis is not mentioned in an input pspec, we can always rewrite to aless efficient program where all mesh axes are mentioned but the caller performsajnp.tile, for example:

@partial(shard_map,mesh=m,in_specs=P('i','j'),out_specs=P('i','j'))deff2(x_block):print(x_block.shape)returnx_blockx=np.arange(12*12).reshape(12,12)x_=jnp.tile(x,(1,mesh.axis_size['j']))# x_ has shape (12, 24)y=f2(x_)# prints (3,12), and f1(x) == f2(x_)

In other words, because each input pspec can mention each mesh axis name zero orone times, rather than having to mention each name exactly once, we can say thatin addition to thejnp.split built into its input,shard_map also has ajnp.tile built into its input, at least logically (though the tiling may notneed to be carried out physically, depending on the arguments’ physical shardinglayout). The tiling to use is not unique; we could also have tiled along thefirst axis, and used the pspecP(('j','i'),None).

Physical data movement is possible on inputs, as each device needs to have acopy of the appropriate data.

Controlling how each output assembled by concatenation, block transposition, and untiling usingout_specs#

Analogously to the input side, each of theout_specs identifies some of thecorresponding output array’s axes with mesh axes by name, representing how theoutput blocks (one for each application of the body function, or equivalentlyone for each physical device) should be assembled back together to form thefinal output value. For example, in both thef1 andf2 examples above theout_specs indicate we should form the final output by concatenating togetherthe block results along both axes, resulting in both cases an arrayy of shape(12,24). (It’s an error if an output shape of the body function, i.e. anoutput block shape, has a rank too small for the concatenation described by thecorresponding output pspec.)

When a mesh axis name is not mentioned in an output pspec, it represents anun-tiling: when the user writes an output pspec which does not mention one ofthe mesh axis names, they promise that the output blocks are equal along thatmesh axis, and so only one block along that axis is used in the output (ratherthan concatenating all the blocks together along that mesh axis). For example,using the same mesh as above:

x=jnp.array([[3.]])z=shard_map(lambda:x,mesh=m,in_specs=(),out_specs=P('i','j'))()print(z)# prints the same as jnp.tile(x, (4, 2))z=shard_map(lambda:x,mesh=m,in_specs=(),out_specs=P('i',None))()print(z)# prints the same as jnp.tile(x, (4, 1)), or just jnp.tile(x, (4,))z=shard_map(lambda:x,mesh=m,in_specs=(),out_specs=P(None,None))()print(z)# prints the same as jnp.tile(x, (1, 1)), or just x

Notice that the body function closing over an array value is equivalent topassing it as an augment with a corresponding input pspec ofP(None,None). Asanother example, following more closely to the other examples above:

@partial(shard_map,mesh=m,in_specs=P('i','j'),out_specs=P('i',None))deff3(x_block):returnjax.lax.psum(x_block,'j')x=np.arange(12*12).reshape(12,12)y3=f3(x)print(y3.shape)# (12,6)

Notice that the result has a second axis size of 6, half the size of the input’ssecond axis. In this case, the un-tile expressed by not mentioning the mesh axisname'j' in the output pspec was safe because of the collectivepsum, whichensures each output block is equal along the corresponding mesh axis. Here aretwo more examples where we vary which mesh axes are mentioned in the outputpspec:

@partial(shard_map,mesh=m,in_specs=P('i','j'),out_specs=P(None,'j'))deff4(x_block):returnjax.lax.psum(x_block,'i')x=np.arange(12*12).reshape(12,12)y4=f4(x)print(y4.shape)# (3,12)@partial(shard_map,mesh=m,in_specs=P('i','j'),out_specs=P(None,None))deff5(x_block):returnjax.lax.psum(x_block,('i','j'))y5=f5(x)print(y5.shape)# (3,6)

On the physical side, not mentioning a mesh axis name in an output pspecassembles anArray from the output device buffers with replicated layout alongthat mesh axis.

There is no runtime check that the output blocks are actually equal along a meshaxis to be un-tiled along, or equivalently that the corresponding physicalbuffers have equal values and thus can be interpreted as a replicated layout fora single logical array. But we can provide a static check mechanism which raisesan error on all potentially-incorrect programs.

Because theout_specs can mention mesh axis names zero or one times, andbecause they can be mentioned in any order, we can say that in addition to thejnp.concatenate built into its output,shard_map also has both an untile anda block transpose built into its output.

Physical data movement is not possible on outputs, no matter the output pspec.Instead,out_specs just encodes how to assemble the block outputs intoArrays, or physically how to interpret the buffers across devices as thephysical layout of a single logicalArray.