def test_unsorted_segment_softmax(self, data, expected_softmax):
    """Verifies variable names and shapes created by a DeepSets network."""

    data = tf.constant(data, dtype=tf.float32)
    segment_ids = tf.constant(self.RECEIVERS, dtype=tf.int32)
    num_segments = tf.constant(sum(self.N_NODE), dtype=tf.int32)

    actual_softmax = modules._unsorted_segment_softmax(
        data, segment_ids, num_segments)

    self.assertAllClose(expected_softmax, actual_softmax.numpy())
Verifies variable names and shapes created by a DeepSets network.
graph_nets/tests_tf2/modules_test.py