code: def _make_model(
      neighbor_config, inputs, keep_rank, weight_dtype=None
  ):
      """Makes a model that exercises the `NeighborFeatures` layer.
  
      The model takes a dictionary of sample features as input and unpacks the
      dictionary into sample features, neighbor features, and neighbor weights.
  
      Args:
        neighbor_config: An instance of `configs.GraphNeighborConfig`.
        inputs: A `tf.keras.Input` or a nested structure of `tf.keras.Input`s.
        keep_rank: Whether to retain the rank of the original input tensors by
          merging the neighborhood size with the batch_size dimension, or add an
          extra neighborhood size dimension.
        weight_dtype: Optional `tf.DType` for weights.
  
      Returns:
        An instance of `tf.keras.Model`.
      """
      # Create inputs for neighbor features and unpack.
      neighbor_features_layer = neighbor_features_lib.NeighborFeatures(
          neighbor_config, weight_dtype=weight_dtype
      )
      (
          sample_features,
          neighbor_features,
          neighbor_weights,
      ) = neighbor_features_layer(inputs, keep_rank=keep_rank)
      return tf.keras.Model(
          inputs=inputs,
          outputs=(
              sample_features,
              neighbor_features,
              neighbor_weights,
          ),
      )
  
file_path: neural-structured-learning/neural_structured_learning/keras/layers/neighbor_features_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import neural_structured_learning.configs
  from neural_structured_learning.keras.layers import neighbor_features
  import numpy
  import tensorflow
  from tensorflow.python.framework import test_util
label_desc: Use image node_id 1 for calling a global function with example usage: _make_model(neighbor_config, inputs, keep_rank, weight_dtype) and returns: tf