The XLA and vanilla TF variants appear to be the same here: https://github.com/tensorflow/tensorflow/blob/2c8d0dca978a246f54c506aae4587dbce5d3bcf0/tensorflow/c

The XLA and vanilla TF variants appear to be the same here: https://github.com/tensorflow/tensorflow/blob/2c8d0dca978a246f54c506aae4587dbce5d3bcf0/tensorflow/compiler/tf2xla/kernels/softmax_op.cc#L50 https://github.com/tensorflow/tensorflow/blob/2c8d0dca978a24... https://github.com/tensorflow/tensorflow/blob/2c8d0dca978a246f54c506aae4587dbce5d3bcf0/tensorflow/core/kernels/softmax_op_functor.h#L67 https://github.com/tensorflow/tensorflow/blob/2c8d0dca978a24... Chalk it up to poetic license? ;-)