Repository: google/jax

Issue:
Undefined variable x
I'm hitting this line in `pxla.py` which doesn't seem to be tested right now:

https://github.com/google/jax/blob/505b47c1fcef45534aa0cc90509ce146f9939c15/jax/interpreters/pxla.py#L80

The variable `x` isn't defined anywhere.


Discussion:
Gooooood catch! I think I know how to write a repro, but if you've got one please share it.
My code is too complicated unfortunately :(
I'm not sure how to make a minimal repro so if you can think of a way I would say go ahead.
Here's a repro when 8 XLA devices are available (like in our favorite colab notebooks):

```python
from jax import pmap
import numpy as onp
a = pmap(lambda x: x)(onp.ones((2, 2)))
b = pmap(pmap(lambda x: x))(a)
```
By the way, as a workaround you can always bounce to the host (calling `onp.asarray`) or to a single device (calling `device_put`). I'm working on a more efficient solution though.

Let me know if you need better workarounds.
Meetings!

Okay, here's something you can paste into a colab cell to fix things:

```python
from jax.interpreters import pxla, xla
from jax.interpreters.pxla import ShardedDeviceArray, shard_arg_handlers, assign_shards_to_replicas

def shard_args(backend, devices, assignments, axis_size, tuple_args, args):
  nargs, nrep = len(args), len(devices)
  buffers = [[None] * nargs for _ in range(nrep)]
  for a, arg in enumerate(args):
    # The shard_arg_handlers allow an extensible set of types to be sharded, but
    # inline handling for ShardedDeviceArray as a special case for performance
    if type(arg) is ShardedDeviceArray:
      if nrep == len(arg.device_buffers):
        # The argument is already prepared for the right number of replicas, so
        # we just ensure that buf[r] is on devices[r] for each replica index r
        for r, buf in enumerate(arg.device_buffers):
          buffers[r][a] = buf if buf.device() == devices[r] else buf.copy_to_device(devices[r])
      else:
        # The argument is prepared for a different number of replicas, so for
        # each of our replica indices we check if there's already a buffer with
        # the correct logical assignment on the correct device, and if not just
        # copy one of them
        prev_assignments = assign_shards_to_replicas(len(arg.device_buffers), axis_size)
        for r in range(nrep):
          candidates = (buf for i, buf in enumerate(arg.device_buffers)
                        if prev_assignments[i] == assignments[r])
          for buf in candidates:
            if buf.device() == devices[r]:
              buffers[r][a] = buf
              break
          else:
            buffers[r][a] = buf.copy_to_device(devices[r])
    else:
      bufs = shard_arg_handlers[type(arg)](arg, devices, assignments, backend=backend)
      for r, buf in enumerate(bufs):
        buffers[r][a] = buf

  if tuple_args:
    buffers = [[xla.make_tuple(bufs, devices[r], backend)]
               for r, bufs in enumerate(buffers)]

  return buffers
pxla.shard_args = shard_args
```

I'll get this merged into master and google asap.

Resolving patch:
diff --git a/jax/interpreters/pxla.py b/jax/interpreters/pxla.py
--- a/jax/interpreters/pxla.py
+++ b/jax/interpreters/pxla.py
@@ -16,7 +16,7 @@
 from __future__ import division
 from __future__ import print_function
 
-from collections import namedtuple
+from collections import namedtuple, defaultdict
 from contextlib import contextmanager
 import itertools as it
 import operator as op
@@ -49,16 +49,16 @@
 def identity(x): return x
 
 def shard_args(backend, devices, assignments, axis_size, tuple_args, args):
-  """Shard an argument data array arg along its leading axis.
+  """Shard each argument data array along its leading axis.
 
   Args:
-    devices: list of Devices of length num_replicas mapping a logical replica
-      index to a physical device.
-    assignments: replica to shard assignment.
-    axis_size: int, size of the axis to be sharded.
+    backend: the platform to be used
+    devices: list of Devices mapping replica index to a physical device.
+    assignments: list of integers with the same length as `devices` mapping
+      replica index to an index along the leading axis (i.e. a shard).
+    axis_size: int, size of the leading axis to be sharded.
     args: a sequence of JaxTypes representing arguments to be sharded along
-      their leading axes (or the leading axess of their leaves in the tuple
-      case) and placed on `devices`.
+      their leading axes and placed on `devices`.
 
   Returns:
     A list of device buffers with the same length as `devices` indexed by
@@ -72,12 +72,28 @@ def shard_args(backend, devices, assignments, axis_size, tuple_args, args):
     # inline handling for ShardedDeviceArray as a special case for performance
     if type(arg) is ShardedDeviceArray:
       if nrep == len(arg.device_buffers):
+        # The argument is already prepared for the right number of replicas, so
+        # we just ensure that buf[r] is on devices[r] for each replica index r
+        # TODO(mattjj): compared to the other case, this logic has less looping
+        # but could incur more device-to-device data movement
         for r, buf in enumerate(arg.device_buffers):
-          buffers[r][a] = (buf if buf.device() == devices[r]
-                           else buf.copy_to_device(devices[r]))
+          buffers[r][a] = buf if buf.device() == devices[r] else buf.copy_to_device(devices[r])
       else:
+        # The argument is prepared for a different number of replicas, so for
+        # each of our replica indices we check if there's already a buffer with
+        # the correct logical assignment on the correct device, and if not just
+        # copy one of them
+        prev_assignments = assign_shards_to_replicas(len(arg.device_buffers), axis_size)
+        candidates = defaultdict(list)
         for r, buf in enumerate(arg.device_buffers):
-          buffers[r][a] = xla.device_put(x[assignments[r]], devices[r], backend=backend)
+          candidates[prev_assignments[r]].append(buf)
+        for r in range(nrep):
+          for buf in candidates[assignments[r]]:
+            if buf.device() == devices[r]:
+              buffers[r][a] = buf
+              break
+          else:
+            buffers[r][a] = buf.copy_to_device(devices[r])
     else:
       bufs = shard_arg_handlers[type(arg)](arg, devices, assignments, backend=backend)
       for r, buf in enumerate(bufs):
@@ -89,6 +105,7 @@ def shard_args(backend, devices, assignments, axis_size, tuple_args, args):
 
   return buffers
 
+
 shard_arg_handlers = {}
 shard_arg_handlers[core.Unit] = \
     lambda x, devices, _, backend=None: [


Test patch: