Remove assert from ShardedDeviceArray staging. (jax-ml#2908)

This would erroneously fail on Cloud TPU because the TPU client has its own buffer type.

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:])