code: def test_compute_sampled_logits(self):
      cs_config = cs_config_builder.build_candidate_sampler_config(
          cs_config_builder.negative_sampler(
              unique=True, algorithm="UNIFORM"
          )
      )
      de_config = test_util.default_de_config(3, cs_config=cs_config)
  
      # Add a few embeddings into knowledge bank.
      de_ops.dynamic_embedding_update(
          ["key1", "key2", "key3"],
          tf.constant(
              [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]
          ),
          de_config,
          "emb",
          service_address=self._kbs_address,
      )
  
      # Sample logits.
      (
          logits,
          labels,
          keys,
          mask,
          weights,
      ) = cs_ops.compute_sampled_logits(
          [["key1", ""], ["key2", "key3"]],
          tf.constant([[2.0, 4.0, 1], [-2.0, -4.0, 1]]),
          3,
          de_config,
          "emb",
          service_address=self._kbs_address,
      )
  
      # Expected results:
      # - Example one returns one positive key {'key2'} and two negative keys
      #   {'key2', 'key3'}.
      # - Example two returns two positive keys {'key2', 'key3'} and one
      #   positive key {'key1'}.
      expected_weights = {
          b"key1": [1, 2, 3],
          b"key2": [4, 5, 6],
          b"key3": [7, 8, 9],
      }
      expected_labels = [
          {b"key1": 1, b"key2": 0, b"key3": 0},
          {b"key1": 0, b"key2": 1, b"key3": 1},
      ]
      # Logit for example one:
      # - 'key1': [2, 4, 1] * [1, 2, 3] = 13
      # - 'key2': [2, 4, 1] * [4, 5, 6] = 34
      # - 'key3': [2, 4, 1] * [7, 8, 9] = 55
      # Logit for example two:
      # - 'key1': [-2, -4, 1] * [1, 2, 3] = -7
      # - 'key2': [-2, -4, 1] * [4, 5, 6] = -22
      # - 'key3': [-2, -4, 1] * [7, 8, 9] = -37
      expected_logits = [
          {b"key1": 13, b"key2": 34, b"key3": 55},
          {b"key1": -7, b"key2": -22, b"key3": -37},
      ]
      # Check keys and weights.
      for b in range(2):
          self.assertEqual(1, mask.numpy()[b])
  
          for key in {b"key1", b"key2", b"key3"}:
              self.assertIn(key, keys.numpy()[b])
          for i in range(3):
              key = keys.numpy()[b][i]
              self.assertAllClose(
                  expected_weights[key], weights.numpy()[b][i]
              )
              self.assertAllClose(
                  expected_labels[b][key], labels.numpy()[b][i]
              )
              self.assertAllClose(
                  expected_logits[b][key], logits.numpy()[b][i]
              )
  
file_path: neural-structured-learning/research/carls/candidate_sampling_ops_test.py
imports: from research.carls import candidate_sampling_ops
  from research.carls import context
  from research.carls import dynamic_embedding_ops
  from research.carls.candidate_sampling import candidate_sampler_config_builder
  from research.carls.testing import test_util
  import tensorflow
label_desc: Use image node_id 4 for calling the CandidateSamplingOpsTest obj's underlying member method code with example usage: obj.test_compute_sampled_logits() without return types