make einsum deterministic, correct.
make einsum deterministic, correct.

Fixes a nondeterministic batch-dimension reordering error that was caused by using a python set collection ordering to fix the final output permutations
tensorflow/probability,google/jax,google/jax,google/jax,tensorflow/probability,google/jax