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: Initializes the layer.
  
  Args:
    num_heads: Number of attention heads (see `MultiHeadAttention` for more
      details on this argument).
    head_size: Size of each attention head.
    num_layers: Number of cross-document attention layers.
    dropout: Dropout probability.
    name: Name of the layer.
    input_noise_stddev: Input Gaussian noise standard deviation.
    **kwargs: keyword arguments.
docstr_tok: ["Initializes","the","layer",".","Args",":","num_heads",":","Number","of","attention","heads","(","see","`","MultiHeadAttention","`","for","more","details","on","this","argument",")",".","head_size",":","Size","of","each","attention","head",".","num_layers",":","Number","of","cross-document","attention","layers",".","dropout",":","Dropout","probability",".","name",":","Name","of","the","layer",".","input_noise_stddev",":","Input","Gaussian","noise","standard","deviation",".","*","*","kwargs",":","keyword","arguments","."]
code: def __init__(
      self,
      num_heads: int,
      head_size: int,
      num_layers: int = 1,
      dropout: float = 0.5,
      name: Optional[str] = None,
      input_noise_stddev: Optional[float] = None,
      **kwargs: Dict[Any, Any],
  ):
      """Initializes the layer.
  
      Args:
        num_heads: Number of attention heads (see `MultiHeadAttention` for more
          details on this argument).
        head_size: Size of each attention head.
        num_layers: Number of cross-document attention layers.
        dropout: Dropout probability.
        name: Name of the layer.
        input_noise_stddev: Input Gaussian noise standard deviation.
        **kwargs: keyword arguments.
      """
      super().__init__(name=name, **kwargs)
      self._num_heads = num_heads
      self._head_size = head_size
      self._num_layers = num_layers
      self._dropout = dropout
      self._input_noise_stddev = input_noise_stddev
  
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 1 to create a new DocumentInteractionAttention object from inherited base classes: tf.keras.layers with example: obj = DocumentInteractionAttention(num_heads, head_size, num_layers, dropout, name, input_noise_stddev)