jax

af63e443 - [sharding_in_types] Check out_avals with mesh context too. This is because users can pass their own shardings to functions like `einsum`, `reshape, `broadcast`, etc