jax

2717daee - Simplifies full to not instantiate intermediate array with default sharding, this significantly reduces overhead when creating sharded arrays in eager mode when using jnp.zeros_like(...)