def get_scores_unstable(self, temperature=1, lhs_nonterminal_bias=None):
    """Same as get_scores but is less numerically stable."""
    # TODO(petershaw): Consider removing this and just using the above method?
    if lhs_nonterminal_bias is None:
      lhs_nonterminal_bias = tf.zeros_like(self.type_lhs_scores)
    # <float>[num_rhs_emb, num_types]
    type_given_rhs = tf.nn.softmax(self.rhs_type_scores / temperature, axis=1)
    # <float>[num_types, num_lhs_emb]
    lhs_given_type = tf.nn.softmax(
        self.type_lhs_scores / temperature + lhs_nonterminal_bias, axis=1)
    # Avoid broadcast to large matrix.
    logp_lhs_given_rhs = tf.math.log(tf.matmul(type_given_rhs, lhs_given_type))
    logp_lhs_given_rhs_t = tf.transpose(logp_lhs_given_rhs)
    return logp_lhs_given_rhs_t
Same as get_scores but is less numerically stable.
language/compgen/csl/model/weighted_model.py