Context

We would like to add torch::nn::functional::one_hot to the C++ API, so that C++ users can easily find the equivalent of Python API torch.nn.functional.one_hot.

Steps

• Add torch::nn::functional::one_hot(...) in torch/csrc/api/include/torch/nn/functional/embedding.h (add this file if it doesn’t exist). The function should have the following signature:

namespace torch {
namespace nn {
namespace functional {

inline Tensor one_hot(
    const Tensor& tensor,
    int64_t num_classes = -1) {
  ...
}

} // namespace functional
} // namespace nn
} // namespace torch

• Add test for torch::nn::functional::one_hot(...) in test/cpp/api/functional.cpp. It can just check whether the function output is as expected for the following cases (expressed in Python):

>>> F.one_hot(torch.arange(0, 5) % 3)
tensor([[1, 0, 0],
        [0, 1, 0],
        [0, 0, 1],
        [1, 0, 0],
        [0, 1, 0]])
>>> F.one_hot(torch.arange(0, 5) % 3, num_classes=5)
tensor([[1, 0, 0, 0, 0],
        [0, 1, 0, 0, 0],
        [0, 0, 1, 0, 0],
        [1, 0, 0, 0, 0],
        [0, 1, 0, 0, 0]])
>>> F.one_hot(torch.arange(0, 6).view(3,2) % 3)
tensor([[[1, 0, 0],
         [0, 1, 0]],
        [[0, 0, 1],
         [1, 0, 0]],
        [[0, 1, 0],
         [0, 0, 1]]])

Helpful Resources

There are quite a few PRs for adding new functionals / new modules for the C++ API (the list of PRs is in #25883), which can serve as great references. Also please ping @yf225 on this issue if you encounter any problems.

How do I claim this feature request?

Please comment in this issue if you are interested in working on it.

cc @yf225