[sharding_in_types][Take 2] Add out_type argument to einsum and dot_general to allow specifying for the output type. Right now, it only accept a NamedSharding but in the future we can allow a polymorphic type of: jax.ShapeDtypeStruct | Sharding | Layout.#743
Open
copybara-service[bot] wants to merge 1 commit intomainfrom test_688540505
+2
