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: Updates global states based on the cluster assignments of the new points.
  
  The number of possible clusters and the shape of the tensors containing
  cluster repesentations are dynamically expanded if any row in the batch no
  longer has an empty cluster (due to the new point joining that cluster).
  
  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 from the output of the
      `call()` method.
    cluster_ids: An int `Tensor` of shape `[batch_size]`. The cluster IDs
      assigned to the new points.
  
  Returns:
    A list of updated states.
code: def update_states_by_cluster_assignment(
      self, next_preprocessed_input, states, cluster_ids
  ):
      """Updates global states based on the cluster assignments of the new points.
  
      The number of possible clusters and the shape of the tensors containing
      cluster repesentations are dynamically expanded if any row in the batch no
      longer has an empty cluster (due to the new point joining that cluster).
  
      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 from the output of the
          `call()` method.
        cluster_ids: An int `Tensor` of shape `[batch_size]`. The cluster IDs
          assigned to the new points.
  
      Returns:
        A list of updated states.
      """
      h_x_next, _ = next_preprocessed_input
      batch_size = h_x_next.shape[0]
      u_aggr, h_aggr, g_h, g_aggr, g_candidates, num_clusters = states
  
      # The tensor indices used to update g_aggr, h_aggr, g_h.
      indices = tf.stack([tf.range(batch_size), cluster_ids], axis=1)
  
      # Non-batch equivalent: g_aggr = g_candidates[cluster_id]
      g_aggr = tf.gather_nd(g_candidates, indices)
      g_candidates = None
  
      # Non-batch equivalent: h_aggr[cluster_id] += h_x_next.
      h_aggr = tf.tensor_scatter_nd_add(h_aggr, indices, h_x_next)
  
      # Non-batch equivalent:
      # g_h[cluster_id] = self.cluster_layer(h_aggr[cluster_id])
      g_h_update = self.cluster_layer(tf.gather_nd(h_aggr, indices))
      g_h = tf.tensor_scatter_nd_update(g_h, indices, g_h_update)
  
      # Dynamically expand the number of clusters by growing the embeddings.
      max_num_clusters = h_aggr.shape[1]
      max_cluster_id = np.max(cluster_ids)
      if max_cluster_id == max_num_clusters - 1:
          h_aggr = tf.concat(
              [
                  h_aggr,
                  tf.zeros([h_aggr.shape[0], 1, h_aggr.shape[-1]]),
              ],
              axis=1,
          )
          g_h = tf.concat(
              [g_h, tf.zeros(shape=[g_h.shape[0], 1, g_h.shape[-1]])],
              axis=1,
          )
  
      for i, k in enumerate(cluster_ids):
          if k == num_clusters[i] - 1:
              num_clusters[i] += 1
  
      states = [u_aggr, h_aggr, g_h, g_aggr, g_candidates, num_clusters]
      return states
  
file_path: neural-structured-learning/research/neural_clustering/models/ncp_base.py
imports: import numpy
  import tensorflow
label_desc: Use image node_id 5 for calling the NCPBase obj's underlying member method code with example usage: obj.update_states_by_cluster_assignment(next_preprocessed_input, states, cluster_ids) and returns: states