Remove assert from ShardedDeviceArray staging. (#2908)

This would erroneously fail on Cloud TPU because the TPU client has its own buffer type.
google · Apr 30, 2020 · 815a92e · 815a92e
1 parent 3216f5c
commit 815a92e
Showing 1 changed file with 0 additions and 1 deletion.
diff --git a/jax/interpreters/pxla.py b/jax/interpreters/pxla.py
@@ -459,7 +459,6 @@ def __init__(self,
     # 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:])