F.softmax should take a dimension on which to perform the softmax (like https://www.tensorflow.org/api_docs/python/tf/nn/softmax)

F.softmax(input, dim=-1)

Currently, the choice of dimension is somewhat strange and inherited from Lua Torch:
1d: dim 0
2d: dim 1
3d: dim 0
4d: dim 1

Like TensorFlow, I think we should also default to the right-most dimension (-1). This would be a breaking change for 3d and 4d inputs, but I think it would be worth it, rather than trying to describe the current behavior.

(We should also change nn.Softmax to take a dimension in the constructor)