def METHOD_NAME(self, hidden_states, word_embeddings_weight):
    if hasattr(word_embeddings_weight, 'infshape'):
        width_mult = word_embeddings_weight.infshape.width_mult()
    else:
        width_mult = 1.0
        if not self.warn_once:
            logging.warning("need to set_shape before use mu-Transfer readout layer")
        self.warn_once = True
    async_tensor_model_parallel_allreduce = parallel_state.get_tensor_model_parallel_world_size() > 1
    output = parallel_lm_logits(
        hidden_states / width_mult,
        word_embeddings_weight,
        self.parallel_output,
        bias=self.bias,
        async_tensor_model_parallel_allreduce=async_tensor_model_parallel_allreduce,
    )
    return output