jax

0f87a2fc - Error out in lax.tile if a dimension is sharded and the rep is non-trivial (i.e. not 1).