code: def unpack_neighbor_features(
      features, neighbor_config, keep_rank=True
  ):
      """Extracts sample features, neighbor features, and neighbor weights.
  
      For example, suppose `features` contains a single sample feature named
      'F0', the batch size is 2, and each sample has 3 neighbors. Then `features`
      might look like the following:
  
      ```
      features = {
          'F0': tf.constant(11.0, shape=[2, 4]),
          'NL_nbr_0_F0': tf.constant(22.0, shape=[2, 4]),
          'NL_nbr_0_weight': tf.constant(0.25, shape=[2, 1]),
          'NL_nbr_1_F0': tf.constant(33.0, shape=[2, 4]),
          'NL_nbr_1_weight': tf.constant(0.75, shape=[2, 1]),
          'NL_nbr_2_F0': tf.constant(44.0, shape=[2, 4]),
          'NL_nbr_2_weight': tf.constant(1.0, shape=[2, 1]),
      },
      ```
  
      where `NL_nbr_<i>_F0` represents the corresponding neighbor features for the
      sample feature 'F0', and `NL_nbr_<i>_weight` represents its neighbor weights.
      The specific values for each key (tensors) in this dictionary are for
      illustrative purposes only. The first dimension of all tensors is the batch
      size.
  
      Example invocation:
  
      ```
      neighbor_config = nsl.configs.make_graph_reg_config(max_neighbors=3)
      sample_features, nbr_features, nbr_weights = nsl.lib.unpack_neighbor_features(
          features, neighbor_config)
      ```
  
      After performing these calls, we would have `sample_features` set to:
  
      ```
      { 'F0': tf.constant(11.0, shape=[2, 4]) },
      ```
  
      `neighbor_features` set to:
  
      ```
      # The key in this dictionary will contain the original sample's feature name.
      # The shape of the corresponding tensor will be 6x4, which is the result of
      # doing an interleaved merge of three 2x4 tensors along axis 0.
      {
        'F0': tf.constant([[22, 22, 22, 22], [33, 33, 33, 33], [44, 44, 44, 44],
                           [22, 22, 22, 22], [33, 33, 33, 33], [44, 44, 44, 44]]),
      },
      ```
      and `neighbor_weights` set to:
  
      ```
      # The shape of this tensor is 6x1, which is the result of doing an
      # interleaved merge of three 2x1 tensors along axis 0.
      tf.constant([[0.25], [0.75], [1.0], [0.25], [0.75], [1.0]])
      ```
  
      Args:
        features: Dictionary of tensors mapping feature names (sample features,
          neighbor features, and neighbor weights) to tensors. For each sample
          feature, all its corresponding neighbor features and neighbor weights must
          be included. All tensors should have a rank that is at least 2, where the
          first dimension is the batch size. The shape of every sample feature
          tensor should be identical to each of its corresponding neighbor feature
          tensors. The shape of each neighbor weight tensor is expected to be `[B,
          1]`, where `B` is the batch size. Neighbor weight tensors cannot be sparse
          tensors.
        neighbor_config: An instance of `nsl.configs.GraphNeighborConfig`.
        keep_rank: Boolean indicating whether to retain the rank from the input or
          to introduce a new dimension for the neighborhood size (axis 1). Defaults
          to `True`.
  
      Returns:
        sample_features: a dictionary mapping feature names to tensors. The shape
          of these tensors remains unchanged from the input.
        neighbor_features: a dictionary mapping feature names to tensors, where
          these feature names are identical to the corresponding feature names in
          `sample_features`. Further, for each feature in this dictionary, the
          resulting tensor represents an interleaved concatenated version of all
          corresponding neighbor feature tensors that exist. So, if the original
          sample feature has a shape `[B, D_1, D_2, ...., D_d]`, then the shape of
          the returned `neighbor_features` will be `[(BxN), D_1, D_2, ..., D_d]` if
          `keep_rank` is `True`, and `[B, N, D_1, D_2, ..., D_d]` if `keep_rank` is
          `False`. If `num_neighbors` is 0, then an empty dictionary is returned.
        neighbor_weights: a tensor containing floating point weights. If `keep_rank`
          is True, `neighbor_weights` will have shape `[(BxN), 1]`. Otherwise, it
          will have shape `[B, N, 1]` This also represents an interleaved
          concatenation of neighbor weight values across all neighbors. The rank of
          this tensor remains unchanged. If `num_neighbors` is 0, then a value of
          `None` is returned.
  
      Raises:
        KeyError: If the input does not contain all corresponding neighbor features
          for every sample feature.
        ValueError: If the tensors of samples and corresponding neighbors don't have
          the same shape.
      """
  
      def check_shape_compatibility(tensors, expected_shape):
          """Checks shape compatibility of the given tensors with `expected_shape`.
  
          Args:
            tensors: List of tensors whose static shapes will be checked for
              compatibility with `expected_shape`.
            expected_shape: Instance of `TensorShape` representing the expected static
              shape of each tensor in `tensors`.
          """
          for tensor in tensors:
              tensor.get_shape().assert_is_compatible_with(
                  expected_shape
              )
  
      # Iterate through the 'features' dictionary to populate sample_features,
      # neighbor_features, and neighbor_weights in one pass.
      sample_features = dict()
      neighbor_features = dict()
      for feature_name, feature_value in features.items():
          # Every value in 'features' is expected to have rank > 1, i.e, 'features'
          # should have been batched to include the extra batch dimension.
          feature_shape = feature_value.get_shape().with_rank_at_least(
              2
          )
  
          if feature_name.startswith(neighbor_config.prefix):
              continue
  
          sample_features[feature_name] = feature_value
  
          # If graph_reg_config.max_neighbors is 0, then neighbor_feature_list will
          # be empty.
          neighbor_feature_list = [
              features[
                  "{}{}_{}".format(
                      neighbor_config.prefix, i, feature_name
                  )
              ]
              for i in range(neighbor_config.max_neighbors)
          ]
  
          # For a given sample feature, aggregate all of its corresponding neighbor
          # features together. Achieve this by doing an interleaved merge of the
          # neighbor feature tensors across all neighbors.
  
          # Populate the 'neighbor_features' dictionary only if there at least one
          # neighbor feature.
          if neighbor_feature_list:
              check_shape_compatibility(
                  neighbor_feature_list, feature_shape
              )
              neighbor_features[feature_name] = _interleave_and_merge(
                  neighbor_feature_list,
                  tf.shape(input=feature_value),
                  keep_rank,
                  is_sparse=isinstance(
                      feature_value, tf.sparse.SparseTensor
                  ),
              )
  
      # If num_neighbors is 0, then neighbor_weights_list will be empty and
      # neighbor_weights will be 'None'.
      neighbor_weights_list = [
          features[
              "{}{}{}".format(
                  neighbor_config.prefix,
                  i,
                  neighbor_config.weight_suffix,
              )
          ]
          for i in range(neighbor_config.max_neighbors)
      ]
  
      # Neighbor weight tensors should have a shape of [B, 1].
      check_shape_compatibility(neighbor_weights_list, [None, 1])
      neighbor_weights = _interleave_and_merge(
          neighbor_weights_list, [-1, 1], keep_rank
      )
  
      return sample_features, neighbor_features, neighbor_weights
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import neural_structured_learning.configs
  import six
  import tensorflow
label_desc: Use image node_id 14 for calling a global function with example usage: unpack_neighbor_features(features, neighbor_config, keep_rank) and returns: sample_features,  neighbor_features,  neighbor_weights