@dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def softmax(x, axis=-1):
  """Softmax of a tensor.