class_docstr: Tests for layers.NeighborFeatures.
class_docstr_tok: ["Tests","for","layers.NeighborFeatures","."]
docstr: Tests creating image neighbors.
docstr_tok: ["Tests","creating","image","neighbors","."]
code: def testDense(self, keep_rank):
      """Tests creating image neighbors."""
      # Make fake 8x8 images.
      batch_size = 4
      image_height = 8
      image_width = 8
      features = {
          "image": np.random.randint(
              0, 256, size=(batch_size, image_height, image_width, 1)
          ).astype(np.uint8),
          "NL_nbr_0_image": np.random.randint(
              0, 256, size=(batch_size, image_height, image_width, 1)
          ).astype(np.uint8),
          "NL_nbr_1_image": np.random.randint(
              0, 256, size=(batch_size, image_height, image_width, 1)
          ).astype(np.uint8),
          "NL_nbr_2_image": np.random.randint(
              0, 256, size=(batch_size, image_height, image_width, 1)
          ).astype(np.uint8),
          "NL_nbr_0_weight": np.random.uniform(
              size=(batch_size, 1)
          ).astype(np.float32),
          "NL_nbr_1_weight": np.random.uniform(
              size=(batch_size, 1)
          ).astype(np.float32),
          "NL_nbr_2_weight": np.random.uniform(
              size=(batch_size, 1)
          ).astype(np.float32),
      }
  
      num_neighbors = 3
      model = _make_model(
          configs.GraphNeighborConfig(max_neighbors=num_neighbors),
          {
              "image": tf.keras.Input(
                  (image_height, image_width, 1),
                  dtype=tf.uint8,
                  name="image",
              ),
          },
          keep_rank,
      )
      samples, neighbors, weights = self.evaluate(model(features))
      samples, neighbors = (samples["image"], neighbors["image"])
      # Check that samples are unchanged.
      self.assertAllEqual(samples, features["image"])
      # Check that neighbors and weights are grouped together for each sample.
      for i in range(batch_size):
          self.assertAllEqual(
              neighbors[(i * num_neighbors) : ((i + 1) * num_neighbors)]
              if keep_rank
              else neighbors[i],
              np.stack(
                  [
                      features["NL_nbr_0_image"][i],
                      features["NL_nbr_1_image"][i],
                      features["NL_nbr_2_image"][i],
                  ]
              ),
          )
          self.assertAllEqual(
              np.split(weights, batch_size)[i]
              if keep_rank
              else weights[i],
              np.stack(
                  [
                      features["NL_nbr_0_weight"][i],
                      features["NL_nbr_1_weight"][i],
                      features["NL_nbr_2_weight"][i],
                  ]
              ),
          )
  
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 the NeighborFeaturesTest obj's underlying member method code with example usage: obj.testDense(keep_rank) without return types