Function softmax_backward_lower

imperative/python/megengine/xla/rules/nn.py:739–743

(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]])

@register_lower_rule("SoftmaxBackward")
def softmax_backward_lower(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]]):
    assert (
        len(args) == 2 and len(ctx.vars_in) == 2 and len(ctx.vars_out) == 1
    ), f"{len(args)}, {len(ctx.vars_in)}, {len(ctx.vars_out)}"
    return softmax_grad(args[0], args[1], ctx.param["axis"])