class_docstr: Neural Clustering Process (NCP) base model implementing one NCP iteration.
  
  The model and notations are based on Algorithm 1 from the NCP paper
  https://arxiv.org/abs/1901.00409. In amortized probabilistic clustering, the
  goal is to learn a parameterized function $ q_{\theta}(c|x) $ to approximate
  the clustering posterior $ p(c|x) $ , where `x` represents the observations
  and `c` represents the cluster labels. NCP does this using a sequential
  point-wise expansion of $ p(c|x) $:
  $$
    p(c|x) = p(c_1|x_{1:N})p(c_2|c_1, x_{1:N})...p(c_N|c_{1:N-1}, x_{1:N})
  $$
  NCP handles a variable number of clusters by dynamically growing neural
  network representations and a variable-input softmax function.
  Permutation invariant representations are used to preserve the permutation
  symmetries within and between clusters.
  
  This class implements the forward pass for one iteration of the NCP model,
  $ q_{\theta}(c_i|c_{1:i-1}, x_{1:N}) $, which assigns a cluster label to
  the `i`-th point given the context of all points and the labels already
  assigned in previous iterations. Within this iteration, it is possible
  for the `i`-th point to join one of the `K` existing clusters, or create a new
  cluster (i.e. joining an empty cluster). Making this decision requires
  evaluating the `K+1` hypotheses, which has been parallelized as much as
  possible. Computing the cluster labels for `N` points requires `N` iterations.
  
  See `NCPWrapper` as an example of how to use the methods in this class.
  
  Attributes:
    assigned_point_layer: A `tf.keras.Model` or `tf.keras.layers.Layer` for the
      `h` function of NCP, which transforms each assigned point `x_i` to `h_i`,
      in order to build representations of existing clusters (`H_k`) from
      assigned points.
    unassigned_point_layer: A `tf.keras.Model` or `tf.keras.layers.Layer` for
      the `u` function of NCP, which transforms each unassigned point `x_i` to
      `u_i`, in order to build the global representation of unassigned points
      (`U`).
    cluster_layer: A `tf.keras.Model` or `tf.keras.layers.Layer` for the `g`
      function of NCP, which transforms the embedding of each cluster `H_k` to
      `g_k`, in order to build the global representation of already clustered
      points (`G`).
    logits_layer: A `tf.keras.Model` or `tf.keras.layers.Layer` for the `f`
      function of NCP (output dim must be 1), which computes the logits of
      possible cluster assignments.
docstr: Computes the logits of adding the new point to all the possible clusters.
  
  Arguments:
    next_preprocessed_input: A list containing the preprocessed tensors
      of the next data point to be clustered: `[h_x_next, u_x_next]`, where
      `h_x_next = h_x[:, next_point]`, and `u_x_next = u_x[:, next_point]`.
        Here, `[h_x, u_x]` is the output of `self.preprocess_inputs(inputs)`,
        and `next_point` is the index of the next point to be clustered.
    states: A list containing the global states:
    - u_aggr: A `Tensor` of shape `[batch_size, u_dim]`. The global
      representation of currently unassigned points (`U`) by aggregating
      `u_x`.
    - h_aggr: A `Tensor` of shape `[batch_size, max(num_clusters), h_dim]`,
      Each slice along the second dimension is the representation of the k-th
      cluster (`H_k`) by aggregating `h_x` assigned to that cluster.
    - g_h: A `Tensor` of shape `[batch_size, max(num_clusters), g_dim]`,
      representing the output of `self.cluster_layer(h_aggr)`.
    - g_aggr: A `Tensor` of shape `[batch_size, g_dim]`. The global
      representation of already clustered data (`G`) by aggregating `g_h`.
    - g_candidates: `None` as part of the input states, but will be updated to
      a `Tensor` of shape `[batch_size, max(num_clusters), g_dim]` in the
      output states, holding the hypothesized global representations after
      adding `x_i` to each of the possible clusters.
    - num_clusters: A `List[int]` (`length=batch_size`) storing the number of
      possible clusters to join for each batch example at the current
      iteration. This includes the empty cluster. If the new point joins the
      empty cluster, it will effectively create a new cluster.
  
  Returns:
    logits: A `Tensor` of shape `[batch_size, max_num_clusters]`. The log
      likelihood of adding the next point to each of the possible clusters.
      For any row `k` in the batch where `num_clusters[k] < max_num_clusters`,
      `logits` will be padded with `-inf` for non-existing clusters as
      `logits[k, n_cluster[k]:] = -np.inf`.
  
    states: A list containing the updated global states.
code: def call(self, next_preprocessed_input, states):
      """Computes the logits of adding the new point to all the possible clusters.
  
      Arguments:
        next_preprocessed_input: A list containing the preprocessed tensors
          of the next data point to be clustered: `[h_x_next, u_x_next]`, where
          `h_x_next = h_x[:, next_point]`, and `u_x_next = u_x[:, next_point]`.
            Here, `[h_x, u_x]` is the output of `self.preprocess_inputs(inputs)`,
            and `next_point` is the index of the next point to be clustered.
        states: A list containing the global states:
        - u_aggr: A `Tensor` of shape `[batch_size, u_dim]`. The global
          representation of currently unassigned points (`U`) by aggregating
          `u_x`.
        - h_aggr: A `Tensor` of shape `[batch_size, max(num_clusters), h_dim]`,
          Each slice along the second dimension is the representation of the k-th
          cluster (`H_k`) by aggregating `h_x` assigned to that cluster.
        - g_h: A `Tensor` of shape `[batch_size, max(num_clusters), g_dim]`,
          representing the output of `self.cluster_layer(h_aggr)`.
        - g_aggr: A `Tensor` of shape `[batch_size, g_dim]`. The global
          representation of already clustered data (`G`) by aggregating `g_h`.
        - g_candidates: `None` as part of the input states, but will be updated to
          a `Tensor` of shape `[batch_size, max(num_clusters), g_dim]` in the
          output states, holding the hypothesized global representations after
          adding `x_i` to each of the possible clusters.
        - num_clusters: A `List[int]` (`length=batch_size`) storing the number of
          possible clusters to join for each batch example at the current
          iteration. This includes the empty cluster. If the new point joins the
          empty cluster, it will effectively create a new cluster.
  
      Returns:
        logits: A `Tensor` of shape `[batch_size, max_num_clusters]`. The log
          likelihood of adding the next point to each of the possible clusters.
          For any row `k` in the batch where `num_clusters[k] < max_num_clusters`,
          `logits` will be padded with `-inf` for non-existing clusters as
          `logits[k, n_cluster[k]:] = -np.inf`.
  
        states: A list containing the updated global states.
      """
  
      h_x_next, u_x_next = next_preprocessed_input
      u_aggr, h_aggr, g_h, g_aggr, g_candidates, num_clusters = states
  
      # Take the current point out of the embedding of unassigned points (U).
      u_aggr -= u_x_next
  
      # Try adding h_x_next to every possible cluster.
      h_candidates = h_aggr + h_x_next[:, tf.newaxis]
      g_candidates = (
          g_aggr[:, tf.newaxis] + self.cluster_layer(h_candidates) - g_h
      )
  
      u_aggr_expanded = tf.broadcast_to(
          u_aggr[:, tf.newaxis],
          [u_aggr.shape[0], g_candidates.shape[1], u_aggr.shape[-1]],
      )
  
      logits = self.logits_layer(
          tf.concat([g_candidates, u_aggr_expanded], axis=-1)
      )
      logits = tf.nn.log_softmax(tf.squeeze(logits, axis=-1), axis=-1)
  
      # Pads logits of non-existing clusters with -inf.
      logits_mask = np.zeros_like(logits)
      for i, k in enumerate(num_clusters):
          logits_mask[i, k:] = -np.inf
      logits_mask = tf.convert_to_tensor(logits_mask)
      logits += logits_mask
  
      states = [u_aggr, h_aggr, g_h, g_aggr, g_candidates, num_clusters]
      return logits, states
  
file_path: neural-structured-learning/research/neural_clustering/models/ncp_base.py
imports: import numpy
  import tensorflow
label_desc: Use image node_id 4 for calling the NCPBase obj's underlying member method code with example usage: obj.call(next_preprocessed_input, states) and returns: logits,  states