Improve one_hot

[zasdfgbnm]

As discussed in #15208, the following improve can be made on torch.nn.functional.ont_hot:

• Make the num_classes parameter int64_t? instead of int64_t num_classes = -1
• Add optional dtype argument
• Add optional dim argument

cc: @zou3519 @vadimkantorov

cc @albanD @mruberry

[deepaks4077]

@zou3519 : Mind if I take a look at this?

[zou3519]

I think the conclusion from @zasdfgbnm on this was that it's very easy to make this API change on the pytorch side but making the JIT support this as well was a little more involved

[zasdfgbnm]

@zou3519 Yes and no. The optional dtype rely on #15154, and the optional int rely on #15234, those were not merged when I wrote #15208. But now this is available, so that part should not be a problem.

[ssnl]

optional dtype is pretty import IMO. returning a long tensor is useless is many cases.

[fmassa]

Constructor functions should natively expose dtype and device. This is currently inferred from the tensor that is passed, but for consistency we should add it as well.

I'm not sure this deserves to be high-pri though, as one_hot is in concept not very different from diag, which also doesn't expose dtype nor device

[ssnl]

@fmassa I would say that this is not a ctor function. However, it should have a dtype arg (similar to how F.softmax has one) because of common precision issue and conversion needs.

[vadimkantorov]

Also in #33046 I propose changing default dtype to bool/byte from long.

[vadimkantorov]

It would also be good if one_hot supported receiving the class index as python scalar, currently only torch.tensor(class_idx) works

[myron]

it would be also good to support 2D and 3D inputs, to create a multi-channel 2D or 3D array (without the need for view/reshape)