Minor formatting fix for `jax.nn.one_hot` docstring

[albanie]

Adds minor code formatting to the jax.nn.one_hot so that it is self-consistent.

[albanie]

Thanks! I'm not sure this is the right change, becuase [0, num_classes) is not valid code, and so formatting it as code might be misleading. The [ and ) here are math symbols, not code symbols: we could maybe format it as math, but that extra markup makes the source version harder to read without really changing the rendered result at all. Given that, maybe just sticking with the original plain text is best. What do you think?

Thanks! You raise a good point. One option is "code-inside-math" like this: [0, num_classes) but I'm also happy to just leave things as they are.