//  [python]
def call(self,
           sequence_data,
           class_index,
           paragraph_mask=None,
           start_positions=None,
           training=False):
    """Implements call().

    Einsum glossary:
    - b: the batch size.
    - l: the sequence length.
    - h: the hidden size, or input width.
    - k: the start/end top n.

    Args:
      sequence_data: The input sequence data of shape
        `(batch_size, seq_length, input_width)`.
      class_index: The class indices of the inputs of shape `(batch_size,)`.
      paragraph_mask: Invalid position mask such as query and special symbols
        (e.g. PAD, SEP, CLS) of shape `(batch_size,)`.
      start_positions: The start positions of each example of shape
        `(batch_size,)`.
      training: Whether or not this is the training phase.

    Returns:
      A dictionary with the keys `start_predictions`, `end_predictions`,
      `start_logits`, `end_logits`.

      If inference, then `start_top_predictions`, `start_top_index`,
      `end_top_predictions`, `end_top_index` are also included.

    """
    paragraph_mask = tf.cast(paragraph_mask, dtype=sequence_data.dtype)
    class_index = tf.reshape(class_index, [-1])

    seq_length = tf.shape(sequence_data)[1]
    start_logits = self.start_logits_dense(sequence_data)
    start_logits = tf.squeeze(start_logits, -1)
    start_predictions, masked_start_logits = _apply_paragraph_mask(
        start_logits, paragraph_mask)

    compute_with_beam_search = not training or start_positions is None

    if compute_with_beam_search:
      # Compute end logits using beam search.
      start_top_predictions, start_top_index = tf.nn.top_k(
          start_predictions, k=self._start_n_top)
      start_index = tf.one_hot(
          start_top_index, depth=seq_length, axis=-1, dtype=tf.float32)
      # start_index: [batch_size, end_n_top, seq_length]

      start_features = tf.einsum('blh,bkl->bkh', sequence_data, start_index)
      start_features = tf.tile(start_features[:, None, :, :],
                               [1, seq_length, 1, 1])
      # start_features: [batch_size, seq_length, end_n_top, input_width]

      end_input = tf.tile(sequence_data[:, :, None],
                          [1, 1, self._start_n_top, 1])
      end_input = tf.concat([end_input, start_features], axis=-1)
      # end_input: [batch_size, seq_length, end_n_top, 2*input_width]
      paragraph_mask = paragraph_mask[:, None, :]
      end_logits = self.end_logits(end_input)

      # Note: this will fail if start_n_top is not >= 1.
      end_logits = tf.transpose(end_logits, [0, 2, 1])
    else:
      start_positions = tf.reshape(start_positions, [-1])
      start_index = tf.one_hot(
          start_positions, depth=seq_length, axis=-1, dtype=tf.float32)
      # start_index: [batch_size, seq_length]

      start_features = tf.einsum('blh,bl->bh', sequence_data, start_index)
      start_features = tf.tile(start_features[:, None, :], [1, seq_length, 1])
      # start_features: [batch_size, seq_length, input_width]

      end_input = tf.concat([sequence_data, start_features],
                            axis=-1)
      # end_input: [batch_size, seq_length, 2*input_width]
      end_logits = self.end_logits(end_input)
    end_predictions, masked_end_logits = _apply_paragraph_mask(
        end_logits, paragraph_mask)

    output_dict = dict(
        start_predictions=start_predictions,
        end_predictions=end_predictions,
        start_logits=masked_start_logits,
        end_logits=masked_end_logits)

    if not training:
      end_top_predictions, end_top_index = tf.nn.top_k(
          end_predictions, k=self._end_n_top)
      end_top_predictions = tf.reshape(
          end_top_predictions,
          [-1, self._start_n_top * self._end_n_top])
      end_top_index = tf.reshape(
          end_top_index,
          [-1, self._start_n_top * self._end_n_top])
      output_dict['start_top_predictions'] = start_top_predictions
      output_dict['start_top_index'] = start_top_index
      output_dict['end_top_predictions'] = end_top_predictions
      output_dict['end_top_index'] = end_top_index

    # get the representation of CLS
    class_index = tf.one_hot(class_index, seq_length, axis=-1, dtype=tf.float32)
    class_feature = tf.einsum('blh,bl->bh', sequence_data, class_index)

    # get the representation of START
    start_p = tf.nn.softmax(masked_start_logits, axis=-1)
    start_feature = tf.einsum('blh,bl->bh', sequence_data, start_p)

    answer_feature = tf.concat([start_feature, class_feature], -1)
    answer_feature = self.answer_logits_inner(answer_feature)
    answer_feature = self.answer_logits_dropout(answer_feature)
    class_logits = self.answer_logits_output(answer_feature)
    class_logits = tf.squeeze(class_logits, -1)
    output_dict['class_logits'] = class_logits
    return output_dict