Expand docs for the :axes option in Nx.gather/3

[Benjamin-Philip]

The documentation for the :axes option inNx.gather/3 can do with some clarification:

• All the examples pertain to collecting subsets. The docs do not explain if and how you can collect scalars.
• The list "nature" of :axes is not explained. All the examples have only a single axis in the list. What is the difference between axes: [0, 1] and axes: [1]?
• The description "The :axes option controls which dimensions the indexes point to" can be cryptic to someone new to the library
• (optional, potentially too verbose) An example of "leading axes" can be helpful to someone new.

[josevalim]

PRs welcome!

[Benjamin-Philip]

I'd be happy to send a PR except that these were issues (at least the first 2) that I encountered myself while reading the docs, i.e. I don't understand the :axes option. If you're okay with going back and forth a bit, so that I can understand the option, I'll send one in a few days. Or else, somebody else could address this?