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: Initializes the global states.
  
  Arguments:
    preprocessed_inputs: a list of preprocessed inputs `[h_x, u_x]`, which is
      the output of `self.preprocess_inputs(inputs)`.
  
  Returns:
    A list containing the initialized global states:
    - u_aggr: A `Tensor` of shape `[batch_size, u_dim]`. The global
      representation of all unassigned points (`U`) by aggregating `u_x`.
    - h_aggr: A zero `Tensor` of shape `[batch_size, 1, h_dim]`. The initial
      representation of an empty cluster with no point assigned to it.
    - g_h: A zero `Tensor` of shape `[batch_size, 1, g_dim]`. This tensor will
      contain the output of `self.cluster_layer(h_aggr)`.
    - g_aggr: A zero `Tensor` of shape `[batch_size, g_dim]`. This tensor will
      contain the global representation of already clustered data (`G`) by
      aggregating `g_h`.
    - g_candidates: Initialized to `None`. This variable will updated to hold
      the hypothesized global representations after adding `x_i` to each of
      the possible clusters.
    - num_clusters: A `List[int]` (`length=batch_size`) initialized to ones
      indicating that there is only one possible cluster (i.e. the empty
      cluster) to join for each batch element at the first iteration.
code: def initialize_states(self, preprocessed_inputs):
      """Initializes the global states.
  
      Arguments:
        preprocessed_inputs: a list of preprocessed inputs `[h_x, u_x]`, which is
          the output of `self.preprocess_inputs(inputs)`.
  
      Returns:
        A list containing the initialized global states:
        - u_aggr: A `Tensor` of shape `[batch_size, u_dim]`. The global
          representation of all unassigned points (`U`) by aggregating `u_x`.
        - h_aggr: A zero `Tensor` of shape `[batch_size, 1, h_dim]`. The initial
          representation of an empty cluster with no point assigned to it.
        - g_h: A zero `Tensor` of shape `[batch_size, 1, g_dim]`. This tensor will
          contain the output of `self.cluster_layer(h_aggr)`.
        - g_aggr: A zero `Tensor` of shape `[batch_size, g_dim]`. This tensor will
          contain the global representation of already clustered data (`G`) by
          aggregating `g_h`.
        - g_candidates: Initialized to `None`. This variable will updated to hold
          the hypothesized global representations after adding `x_i` to each of
          the possible clusters.
        - num_clusters: A `List[int]` (`length=batch_size`) initialized to ones
          indicating that there is only one possible cluster (i.e. the empty
          cluster) to join for each batch element at the first iteration.
      """
  
      h_x, u_x = preprocessed_inputs
  
      batch_size = u_x.shape[0]
      h_dim = h_x.shape[-1]
      g_dim = self.cluster_layer.compute_output_shape(h_x.shape)[-1]
  
      u_aggr = tf.reduce_sum(u_x, axis=1)
      h_aggr = tf.zeros([batch_size, 1, h_dim])
      g_h = tf.zeros([batch_size, 1, g_dim])
      g_aggr = tf.zeros([batch_size, g_dim])
      g_candidates = None
  
      num_clusters = [1] * batch_size
  
      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 3 for calling the NCPBase obj's underlying member method code with example usage: obj.initialize_states(preprocessed_inputs) and returns: states