Function softmax_grad

imperative/python/megengine/xla/rules/nn.py:722–727

(y: HLOTensor, dy: HLOTensor, axis: int = -1)

def softmax_grad(y: HLOTensor, dy: HLOTensor, axis: int = -1):
    assert isinstance(axis, int), f"axis should be int, but get {axis}({type(axis)})"
    ydy = y * dy
    ydy_sum = ydy.sum(axis=axis, keepdims=True)
    dx = ydy - y * ydy_sum
    return dx


@register_lower_rule(mops.Softmax)