class_docstr: Output layer to restore listwise output shape.
  
  This layer maps the flattened 1D logits or 2D tensor of shape
  [batch_size * list_size, 1] back to 2D of shape [batch_size, list_size] and
  mask the invalid terms to be a defined large negative value.
  
  Example usage:
  
  ```python
  flattened_logits = [1, 0.5, 2, 0, -1, 0]
  mask = [[True, True, False], [True, False, False]]
  logits = RestoreList()(inputs=(flattened_logits, mask))
  logits.numpy()
  # Returns: [[1, 0.5, log(1e-10)], [0, log(1e-10), log(1e-10)]]
  ```
  
  This layer works also for 2D `flattened_logits` like
  [[1], [0.5], [2], [0], [-1], [0]].
  
  When `by_scatter=True`, an nd_indices will be generated using `mask` in the
  same way as `FlattenList`. All values in the `flattened_logits` will be used
  and repeated entries will be averaged.
  
  ```python
  flattened_logits = [1, 0.5, 2, 0, -1, 0]
  mask = [[True, True, False], [True, False, False]]
  logits = RestoreList(by_scatter=True)((flattened_logits, mask))
  logits.numpy()
  # Returns: [[1.5, 0.5, log(1e-10)], [-1/3, log(1e-10), log(1e-10)]]
  ```
  
  This is because the flattened_logits are treated as circularly padded entries.
  The [1st, 3rd] values [1, 2] are counted to logits[0, 0]. The [4th, 5th, 6th]
  values [0, -1, 0] are counted to logits[1, 0]. Note that We use different
  values for those repeated entries, while they are likely the same in practice.
docstr: Restores listwise shape of flattened_logits.
  
  Args:
    inputs:  A tuple of (flattened_logits, list_mask), which are described
      below.
    * `flattened_logits`: A `Tensor` of predicted logits for each pair of
      query and documents, 1D tensor of shape [batch_size * list_size] or 2D
      tensor of shape [batch_size * list_size, 1].
    * `list_mask`: A boolean `Tensor` of shape [batch_size, list_size] to mask
      out the invalid examples.
  
  Returns:
    A `Tensor` of shape [batch_size, list_size].
  
  Raises:
    ValueError: If `flattened_logits` is not of shape [batch_size * list_size]
      or [batch_size * list_size, 1].
docstr_tok: ["Restores","listwise","shape","of","flattened_logits",".","Args",":","inputs",":","A","tuple","of","(","flattened_logits",",","list_mask",")",",","which","are","described","below",".","*","`","flattened_logits","`",":","A","`","Tensor","`","of","predicted","logits","for","each","pair","of","query","and","documents",",","1D","tensor","of","shape","[","batch_size","*","list_size","]","or","2D","tensor","of","shape","[","batch_size","*","list_size",",","1","]",".","*","`","list_mask","`",":","A","boolean","`","Tensor","`","of","shape","[","batch_size",",","list_size","]","to","mask","out","the","invalid","examples",".","Returns",":","A","`","Tensor","`","of","shape","[","batch_size",",","list_size","]",".","Raises",":","ValueError",":","If","`","flattened_logits","`","is","not","of","shape","[","batch_size","*","list_size","]","or","[","batch_size","*","list_size",",","1","]","."]
code: def call(self, inputs: Tuple[tf.Tensor, tf.Tensor]) -> tf.Tensor:
      """Restores listwise shape of flattened_logits.
  
      Args:
        inputs:  A tuple of (flattened_logits, list_mask), which are described
          below.
        * `flattened_logits`: A `Tensor` of predicted logits for each pair of
          query and documents, 1D tensor of shape [batch_size * list_size] or 2D
          tensor of shape [batch_size * list_size, 1].
        * `list_mask`: A boolean `Tensor` of shape [batch_size, list_size] to mask
          out the invalid examples.
  
      Returns:
        A `Tensor` of shape [batch_size, list_size].
  
      Raises:
        ValueError: If `flattened_logits` is not of shape [batch_size * list_size]
          or [batch_size * list_size, 1].
      """
      flattened_logits, list_mask = inputs
      try:
          logits = tf.reshape(
              flattened_logits, shape=tf.shape(list_mask)
          )
      except:
          raise ValueError(
              "`flattened_logits` needs to be either "
              "1D of [batch_size * list_size] or "
              "2D of [batch_size * list_size, 1]."
          )
      if self._by_scatter:
          nd_indices, _ = utils.padded_nd_indices(is_valid=list_mask)
          counts = tf.scatter_nd(
              nd_indices, tf.ones_like(logits), tf.shape(list_mask)
          )
          logits = tf.scatter_nd(
              nd_indices, logits, tf.shape(list_mask)
          )
          return tf.where(
              tf.math.greater(counts, 0.0),
              logits / counts,
              tf.math.log(_EPSILON),
          )
      else:
          return tf.where(list_mask, logits, tf.math.log(_EPSILON))
  
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 2 for calling the RestoreList obj's underlying member method code with example usage: obj.call(inputs) and returns: tf, tf