# Copyright (c) Megvii, Inc. and its affiliates.

import math

import megengine.functional as F
import megengine.module as M


class LogitsFullyConnected(M.Module):
    """single fully connected layer, mapping embedding to logits with normalized weight
    """

    def __init__(self, num_class, feature_dim):
        super().__init__()
        fc = M.Linear(feature_dim, num_class, bias=False)
        self.weight = fc.weight
        M.init.msra_uniform_(self.weight, a=math.sqrt(5))

    def forward(self, embedding):
        w = 
F.normalize(self.weight, axis=1)