Add pxla.ShardedDeviceArray to jaxval_zeros_likers
This is a one-line fix at https://github.com/google/jax/blob/master/jax/lax/lax.py#L1503

Just adding this as an issue as a reminder for @mattjj 
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
index 2e05b0b2b..12f5ba849 100644
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -1498,9 +1498,10 @@ def zeros_like_array(x):
   return full_like(x, 0)
 
 for t in itertools.chain(dtypes.python_scalar_dtypes.keys(), array_types,
-                         [xla.DeviceArray]):
+                         [xla.DeviceArray, pxla.ShardedDeviceArray]):
   ad_util.jaxval_adders[t] = add
 ad_util.jaxval_zeros_likers[xla.DeviceArray] = zeros_like_array
+ad_util.jaxval_zeros_likers[pxla.ShardedDeviceArray] = zeros_like_array
 
 
 ### primitives
['jax/lax/lax.py']