self.lookup_table = self.add_weight(
        "lookup_table",
        shape=[self.n_token, self.d_embed],
        initializer=self.initializer,
        dtype=self.dtype)