Temporarily make ShardedDeviceArray.__init__ optionally accept old si…

…gnature.

This allows us to incrementally update ShardedDeviceArray creators to the new constructor introduced in 07571ae.

jax/interpreters/pxla.py

@@ -450,6 +450,15 @@ def __init__(self,
                sharding_spec: ShardingSpec,
                device_buffers: List[xb.xla_client._xla.PyLocalBuffer],
                indices: Optional[Tuple[Index, ...]] = None):
+    # TODO(skye): this is temporary staging while we switch users over to
+    # providing sharding_spec. It assumes that any pre-existing callers are
+    # creating pmap-style ShardedDeviceArrays.
+    if device_buffers is None:
+      assert isinstance(sharding_spec[0], xb.xla_client._xla.PyLocalBuffer)
+      device_buffers = sharding_spec
+      sharding_spec = _pmap_sharding_spec(aval.shape[0], aval.shape[0],
+                                          aval.shape[1:])
+
     # TODO(skye): assert invariants. Keep performance in mind though.
     if indices is None:
       indices = spec_to_indices(aval.shape, sharding_spec)