flax

8e0228e7 - avoid mixing `einsum_dot_general` and `einsum` argument by specifying them explicitly in the caller.