megengine.functional.nn.indexing_one_hot

indexing_one_hot(src, index, axis=1, keepdims=False)[source]

One-hot indexing for some axes.

Parameters

  • src (Tensor) – input tensor.

  • index (Tensor) – index tensor.

  • axis (int) – axis on src for which values in index index. Default: 1

  • keepdims – whether not to remove the axis in result. Default: False

Examples

>>> src = Tensor([[1.0, 2.0]])
>>> index = Tensor([0])
>>> val = F.indexing_one_hot(src, index)
>>> val.numpy()
array([1.], dtype=float32)
Return type

Tensor

previous

megengine.functional.nn.one_hot

next

megengine.functional.nn.embedding