Fix crash in softmax-xent when some input dimensions are 1.

Before, tf.nn.softmax_cross_entropy_with_logits would fail a CHECK if one input tensor had shape (1, 1) and the other did not.

In particular, the call to ToIndexArray<2> here https://github.com/tensorflow/tensorflow/blob/1f3da84a89702d3b4f234ee83762d738caffe098/tensorflow/core/kernels/xent_op.cc#L99 would fail, since the call assumed the array had two dimensions. If both dimensions were 1, BCast would merge the two dimensions into a single dimension. Passing fewer_dims_optimization=false stops this optimization

PiperOrigin-RevId: 384844496
Change-Id: Ifb02dc74964132c3ed3f3bc98b0858dbe4e258b7