Model partitioning when looping jnp.einsum

[JohnG-1qbit]

Hi,

I have a function that applies jnp.einsum to an array in a loop, such that the array is updated after each operation. I'm interested in parallelizing this function to take advantage of model partitioning. My goal is to use model parallelism over multiple CPU devices, and the only function I'm aware of that does this is xmap (pjit model partitioning is not yet implemented for CPUs). Here is an example of what I'm trying to do:

DEVICE_COUNT: int = 4
import os
os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=' + str(DEVICE_COUNT)
import jax
import jax.numpy as jnp
import numpy as onp
from jax.experimental.maps import xmap, mesh
from jax.lax import scan
from jax import partial

def seq_einsum(n_iterations, x, y):
    def _einsum(y, _):
        return jnp.einsum('{i,j},{j, k}->{i, k}', x, y), _
    return scan(_einsum, y, xs=jnp.ones(n_iterations))[0]

if __name__ == '__main__':
    n_iter = 100
    x = jnp.ones((4, 4))
    y = jnp.ones((4, 4))
    in_axes = [['i', 'j', ...], ['j', 'k', ...]]
    out_axes = ['i', 'k', ...]
    _seq_einsum = partial(seq_einsum, n_iter)
    devices = onp.array(jax.local_devices()).reshape((2, 2))
    with mesh(devices, ('j', 'k')):
        z = xmap(_seq_einsum, in_axes, out_axes)(x, y)

However, one can already tell this will not work, because the named axis for y will become i, k instead of j, k after the first call of _einsum, then the next _einsum call cannot be executed. More precisely, I get this error:

File "/.../python3.9/site-packages/jax/interpreters/batching.py", line 73, in _match_axes
    raise ValueError(msg)
ValueError: vmap has mapped output (axis_name=<UniqueResource j 1>) but out_axes is None

The code does run when the xmap is applied to the interior function _einsum rather than the scanned version, but I have seen this comes with a hefty overhead - and my goal is to take advantage of the speedup that comes with compiling the entire scan together. I have considered renaming the axes of y after each _einsum call, but I could not figure out how to do that in the documentation - and it seems to not be in the spirit of xmap. I would appreciate any comments or suggestions!

[sharadmv]

I believe this is an instance of the bug found in issue #7063. It should be fixed in #7206 .