jax

bf5d1c00 - Mark `out_sharding` as a keyword argument in the `pyi` files too for `jnp.ravel`, `jnp.reshape` and `jnp.dot`