def call(self, wordpiece_encodings, application_span_begin,
           application_span_end, application_rule_idx):
    """Computes scores for a batch of anchored rule applications.

    Args:
      wordpiece_encodings: <float>[batch_size, max_num_wordpieces, bert_dims]
      application_span_begin: <int>[batch_size, max_num_applications]
      application_span_end: <int>[batch_size, max_num_applications]
      application_rule_idx: <int>[batch_size, max_num_applications]

    Returns:
      application_scores: <float>[batch_size, max_num_applications]
    """
    # <float>[batch_size, max_num_applications, bert_dims]
    span_begin_encoding = tf.gather(
        wordpiece_encodings, application_span_begin, batch_dims=1)
    span_end_encoding = tf.gather(
        wordpiece_encodings, application_span_end, batch_dims=1)
    # <float>[batch_size, max_num_applications, bert_dims * 2]
    span_encodings = tf.concat([span_begin_encoding, span_end_encoding], axis=2)
    # <float>[batch_size, max_num_applications, model_dims]
    span_encodings_ffn = self.feed_forward(span_encodings)
    # <float>[batch_size, max_num_applications, 1, model_dims]
    span_encodings_ffn = tf.expand_dims(span_encodings_ffn, 2)

    # <float>[batch_size, max_num_applications, model_dims]
    application_rule_embedddings = self.rule_embeddings(application_rule_idx)
    # <float>[batch_size, max_num_applications, model_dims, 1]
    application_rule_embedddings = tf.expand_dims(application_rule_embedddings,
                                                  3)

    # <float>[batch_size, max_num_applications, 1, 1]
    application_scores = tf.matmul(span_encodings_ffn,
                                   application_rule_embedddings)
    # <float>[batch_size, max_num_applications]
    application_scores = tf.squeeze(application_scores, [2, 3])

    # <float>[batch_size, max_num_applications, 1]
    span_scores = self.span_feed_forward(span_encodings)
    # <float>[batch_size, max_num_applications]
    span_scores = tf.squeeze(span_scores, [2])

    return application_scores + span_scores
Computes scores for a batch of anchored rule applications.

    Args:
      wordpiece_encodings: <float>[batch_size, max_num_wordpieces, bert_dims]
      application_span_begin: <int>[batch_size, max_num_applications]
      application_span_end: <int>[batch_size, max_num_applications]
      application_rule_idx: <int>[batch_size, max_num_applications]

    Returns:
      application_scores: <float>[batch_size, max_num_applications]
    
language/compgen/nqg/model/parser/nqg_model.py