NCE-Loss layer under subword-units input.

def nce_loss_subwords(
        data, label, label_mask, label_weight, embed_weight, vocab_size, num_hidden):
    """NCE-Loss layer under subword-units input.
    """
    # get subword-units embedding.
    label_units_embed = mx.sym.Embedding(data=label,
                                         input_dim=vocab_size,
                                         weight=embed_weight,
                                         output_dim=num_hidden)
    # get valid subword-units embedding with the help of label_mask
    # it's achieved by multiplying zeros to useless units in order to handle variable-length input.
    label_units_embed = mx.sym.broadcast_mul(lhs=label_units_embed,
                                             rhs=label_mask,
                                             name='label_units_embed')
    # sum over them to get label word embedding.
    label_embed = mx.sym.sum(label_units_embed, axis=2, name='label_embed')

    # by boardcast_mul and sum you can get prediction scores in all label_embed inputs,
    # which is easy to feed into LogisticRegressionOutput and make your code more concise.
    data = mx.sym.Reshape(data=data, shape=(-1, 1, num_hidden))
    pred = mx.sym.broadcast_mul(data, label_embed)
    pred = mx.sym.sum(data=pred, axis=2)

    return mx.sym.LogisticRegressionOutput(data=pred,
                                           label=label_weight)