change the einsum_string to "1" to ensure the unsharded maybe not a good idea here, since it would disconnect the relation of softmax_axis between input and output when print the vlog for debug:

"ab1" --> "abc", later developer might not realize that axis "1" and "c" are the same axis in two tensors.

it would be better that:
first: out_dims_mapping[axis] = -1;
then: keep the input_axes unchanged.