Function _get_softmax_axis

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

(ndim: int)

Source from the content-addressed store, hash-verified

1014
1015
1016	def _get_softmax_axis(ndim: int) -> int:
1017	if ndim in (0, 1, 3):
1018	return 0
1019	return 1
1020
1021
1022	def softmax(inp: Tensor, axis: Optional[int] = None) -> Tensor:

softmaxFunction · 0.85

no outgoing calls

no test coverage detected