class_docstr: Cross Document Interaction Attention layer.
  
  This layer implements the cross-document attention described in
  [Pasumarthi et al, 2020][pasumarthi2020].
  
  This layer comprises of several layers of Multi-Headed Attention (MHA)
  applied over the list of documents to attend over itself, using a mask to
  specify valid documents. Optionally, the user can specify the `topk` documents
  as the listwise context that is used to compute the attention per document. If
  `topk` is None, all the documents are used as listwise context to compute the
  attention.
  
  MHA uses scaled dot product attention, with residual connection and layer
  normalization as follows. This transformation is applied for `num_layers`
  times:
  
  ```
  h_i := LayerNorm_i(h_{i-1} + MHA_i(h_{i-1}), TopK(h_{i-1}; k))
  ```
  
  Example usage:
  
  ```python
  # Batch size = 2, list_size = 3.
  inputs =  [[[1., 1.], [1., 0.], [1., 1.]], [[0., 0.], [0., 0.], [0., 0.]]]
  list_mask = [[True, True, False], [True, False, False]]
  dia_layer = DocumentInteractionAttention(
      num_heads=1, head_size=64, num_layers=1)
  dia_output = dia_layer(
      inputs=(inputs, list_mask),
      training=False)
  ```
  
  References:
    - [Permutation Equivariant Document Interaction Network for Neural Learning
       to Rank, Pasumarthi et al, 2020][pasumarthi2020]
  
  [pasumarthi2020]: http://research.google/pubs/pub49364/
docstr: Calls the document interaction layer to apply cross-document attention.
  
  Args:
    inputs: A tuple of `example_inputs`, `list_mask`.
      `example_inputs`: A tensor of shape [batch_size, list_size,
        feature_dims].
      `list_mask`: A boolean tensor of shape [batch_size, list_size], which is
        True for a valid example and False for invalid one. If this is `None`,
        then all examples are treated as valid.
    training: Whether in training or inference mode.
  
  Returns:
    A tensor of shape [batch_size, list_size, head_size].
docstr_tok: ["Calls","the","document","interaction","layer","to","apply","cross-document","attention",".","Args",":","inputs",":","A","tuple","of","`","example_inputs","`",",","`","list_mask","`",".","`","example_inputs","`",":","A","tensor","of","shape","[","batch_size",",","list_size",",","feature_dims","]",".","`","list_mask","`",":","A","boolean","tensor","of","shape","[","batch_size",",","list_size","]",",","which","is","True","for","a","valid","example","and","False","for","invalid","one",".","If","this","is","`","None","`",",","then","all","examples","are","treated","as","valid",".","training",":","Whether","in","training","or","inference","mode",".","Returns",":","A","tensor","of","shape","[","batch_size",",","list_size",",","head_size","]","."]
code: def call(self, inputs: tf.Tensor, training: bool = True) -> tf.Tensor:
      """Calls the document interaction layer to apply cross-document attention.
  
      Args:
        inputs: A tuple of `example_inputs`, `list_mask`.
          `example_inputs`: A tensor of shape [batch_size, list_size,
            feature_dims].
          `list_mask`: A boolean tensor of shape [batch_size, list_size], which is
            True for a valid example and False for invalid one. If this is `None`,
            then all examples are treated as valid.
        training: Whether in training or inference mode.
  
      Returns:
        A tensor of shape [batch_size, list_size, head_size].
      """
      example_inputs, list_mask = inputs
      batch_size = tf.shape(example_inputs)[0]
      list_size = tf.shape(example_inputs)[1]
      if list_mask is None:
          list_mask = tf.ones(
              shape=(batch_size, list_size), dtype=tf.bool
          )
      x = self._input_projection(example_inputs, training=training)
      if self._input_noise:
          x = self._input_noise(x, training=training)
  
      list_mask = tf.cast(list_mask, dtype=tf.int32)
      attention_mask = SelfAttentionMask()([list_mask, list_mask])
  
      for (
          attention_layer,
          dropout_layer,
          norm_layer,
      ) in self._attention_layers:
          output = attention_layer(
              query=x,
              value=x,
              attention_mask=attention_mask,
              training=training,
          )
          output = dropout_layer(output, training=training)
          # Applying residual network here, similar to logic in Transformer.
          x = norm_layer(output + x, training=training)
  
      return x
  
imports: import math
  from typing import Any, Callable, Dict, List, Optional, Tuple
  import tensorflow
  from tensorflow_ranking.python import utils
label_desc: Use image node_id 3 for calling the DocumentInteractionAttention obj's underlying member method code with example usage: obj.call(inputs, training) and returns: x