jax

f41a32c6 - lax.full: add sharding argument