oneflow.nn.functional.one_hot

oneflow.nn.functional.one_hot(input, num_classes=- 1, on_value=1, off_value=0)

This operator generates a onehot Tensor from input Tensor.

If input Tensor’s rank is N, the corresponding onehot Tensor’s rank is N+1.

Parameters
  • input (Tensor) – The input Tensor.

  • num_classes (int) – The length of onehot Tensor.

  • on_value (Union[int, float], optional) – The fill value when x[i] == i. Defaults to 1.

  • off_value (Union[int, float], optional) – The fill value when x[i] != i. Defaults to 0.

Note

The data type of input tensor should be int32 or int64.

Returns

oneflow.Tensor.

For example:

>>> import oneflow as flow
>>> import numpy as np

>>> input=flow.tensor(np.array([0, 3, 1, 2]).astype(np.int64), dtype=flow.int64)
>>> out = flow.nn.functional.one_hot(input, num_classes=5)
>>> out
tensor([[1, 0, 0, 0, 0],
        [0, 0, 0, 1, 0],
        [0, 1, 0, 0, 0],
        [0, 0, 1, 0, 0]], dtype=oneflow.int64)