Function softmax_lower

imperative/python/megengine/xla/rules/nn.py:731–735

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

@register_lower_rule(mops.Softmax)
def softmax_lower(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]]):
    assert (
        len(args) == 1 and len(ctx.vars_in) == 1 and len(ctx.vars_out) == 1
    ), f"{len(args)}, {len(ctx.vars_in)}, {len(ctx.vars_out)}"
    return softmax(args[0], ctx.op.axis)


@register_lower_rule("SoftmaxBackward")