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

Function softmax

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

r"""Applies a :math:`\text{softmax}(x)` function. See :class:`~.module.Softmax` for more details. Examples: >>> import numpy as np >>> x = Tensor(np.arange(-5, 5, dtype=np.float32)).reshape(2,5) >>> out = F.softmax(x) >>> out.numpy().round(decimals=4)

(inp: Tensor, axis: Optional[int] = None)

Source from the content-addressed store, hash-verified

1020
1021
1022def softmax(inp: Tensor, axis: Optional[int] = None) -> Tensor:
1023 r"""Applies a :math:`\text{softmax}(x)` function.
1024
1025 See :class:`~.module.Softmax` for more details.
1026
1027 Examples:
1028 >>> import numpy as np
1029 >>> x = Tensor(np.arange(-5, 5, dtype=np.float32)).reshape(2,5)
1030 >>> out = F.softmax(x)
1031 >>> out.numpy().round(decimals=4)
1032 array([[0.0117, 0.0317, 0.0861, 0.2341, 0.6364],
1033 [0.0117, 0.0317, 0.0861, 0.2341, 0.6364]], dtype=float32)
1034 """
1035 if axis is None:
1036 axis = _get_softmax_axis(len(inp.shape))
1037 if isinstance(axis, list):
1038 offset = inp.max(axis=axis, keepdims=True).detach()
1039 cached = exp(inp - offset)
1040 down = sum(cached, axis=axis, keepdims=True)
1041 return cached / down
1042 else:
1043 op = builtin.Softmax(axis=axis,)
1044 (output,) = apply(op, inp)
1045 return output
1046
1047
1048def instance_norm(

Callers

nothing calls this directly

Calls 6

_get_softmax_axisFunction · 0.85
sumFunction · 0.85
expFunction · 0.70
applyFunction · 0.50
detachMethod · 0.45
maxMethod · 0.45

Tested by

no test coverage detected