class_docstr: Tests for layers.NeighborFeatures.
class_docstr_tok: ["Tests","for","layers.NeighborFeatures","."]
docstr: Tests the layer with a variable number of neighbors.
docstr_tok: ["Tests","the","layer","with","a","variable","number","of","neighbors","."]
code: def testSparse(self, keep_rank):
      """Tests the layer with a variable number of neighbors."""
      batch_size = 4
      input_size = 2
      features = {
          "input": tf.sparse.from_dense(
              np.random.normal(size=(batch_size, input_size))
          ),
          # Every sample but the last has 1 neighbor.
          "NL_nbr_0_input": tf.RaggedTensor.from_row_starts(
              values=np.random.normal(
                  size=(batch_size - 1) * input_size
              ),
              row_starts=[0, 2, 4, 6],
          ).to_sparse(),
          "NL_nbr_0_weight": np.expand_dims(
              np.array([0.9, 0.3, 0.6, 0.0]), -1
          ),
          # Only the 1st and 3rd sample have a second neighbor.
          "NL_nbr_1_input": tf.RaggedTensor.from_row_starts(
              values=np.random.normal(
                  size=(batch_size - 2) * input_size
              ),
              row_starts=[0, 2, 2, 4],
          ).to_sparse(),
          "NL_nbr_1_weight": np.expand_dims(
              np.array([0.25, 0.0, 0.75, 0.0]), -1
          ),
      }
  
      model = _make_model(
          configs.GraphNeighborConfig(max_neighbors=2),
          {"input": tf.keras.Input(input_size, dtype=tf.float64)},
          keep_rank,
          tf.float64,
      )
      samples, neighbors, weights = self.evaluate(model(features))
      # Check that samples are unchanged.
      self.assertAllClose(
          samples["input"].values,
          self.evaluate(features["input"].values),
      )
      # Check that weights are grouped together and have the right shape.
      self.assertAllClose(
          weights,
          np.array([0.9, 0.25, 0.3, 0.0, 0.6, 0.75, 0.0, 0.0]).reshape(
              (batch_size * 2, 1) if keep_rank else (batch_size, 2, 1)
          ),
      )
      # Check that neighbors are grouped together.
      dense_neighbors = self.evaluate(
          tf.sparse.to_dense(neighbors["input"], -1.0)
      )
      neighbor0 = self.evaluate(
          tf.sparse.to_dense(features["NL_nbr_0_input"], -1)
      )
      neighbor1 = self.evaluate(
          tf.sparse.to_dense(features["NL_nbr_1_input"], -1)
      )
      for i in range(batch_size):
          actual = (
              np.split(dense_neighbors, batch_size)[i]
              if keep_rank
              else dense_neighbors[i]
          )
          self.assertAllEqual(
              actual, np.stack([neighbor0[i], neighbor1[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 2 for calling the NeighborFeaturesTest obj's underlying member method code with example usage: obj.testSparse(keep_rank) without return types