jax

9deb7e3d - [sharding_in_types] `physical_aval` should set the correct sharding on `ShapedArray` so that lowering and compilation don't crash