MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / indexing_one_hot

Function indexing_one_hot

imperative/python/megengine/functional/nn.py:1696–1720  ·  view source on GitHub ↗

r"""One-hot indexing for some axes. Args: src: input tensor. index: index tensor. axis: 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]]

(
    src: Tensor, index: Tensor, axis: int = 1, keepdims=False
)

Source from the content-addressed store, hash-verified

1694
1695
1696def indexing_one_hot(
1697 src: Tensor, index: Tensor, axis: int = 1, keepdims=False
1698) -> Tensor:
1699 r"""One-hot indexing for some axes.
1700
1701 Args:
1702 src: input tensor.
1703 index: index tensor.
1704 axis: axis on src for which values in index index. Default: 1
1705 keepdims: whether not to remove the axis in result. Default: False
1706
1707 Examples:
1708 >>> src = Tensor([[1.0, 2.0]])
1709 >>> index = Tensor([0])
1710 >>> val = F.indexing_one_hot(src, index)
1711 >>> val.numpy()
1712 array([1.], dtype=float32)
1713 """
1714 assert isinstance(src, Tensor), "src must be of Tensor type"
1715 op = builtin.IndexingOneHot(axis=axis, ndim=src.ndim)
1716 index = convert_single_value(index, dtype="int32", device=src.device)
1717 (result,) = apply(op, src, index)
1718 if not keepdims:
1719 result = squeeze(result, axis)
1720 return result
1721
1722
1723def sliding_window(

Callers 1

cross_entropyFunction · 0.85

Calls 3

convert_single_valueFunction · 0.85
squeezeFunction · 0.85
applyFunction · 0.50

Tested by

no test coverage detected