import megengine as mge
import megengine.functional as F
from megengine.core import Tensor

def softmax_loss(pred, label, ignore_label=-1):
    max_pred = F.zero_grad(pred.max(axis=1, keepdims=True))
    pred -= max_pred
    log_prob = pred - F.log(F.exp(pred).sum(axis=1, keepdims=True))
    mask = 1 - F.equal(label, ignore_label)
    vlabel = label * mask
    loss = -(
F.indexing_one_hot(log_prob, vlabel, 1)