code: def test_compute_sampled_logits_grad(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,
      )
  
      # A simple one layer NN model.
      # Input data: x = [[1, 2], [3, 4]].
      # Weights from input to logit output layer: W = [[1, 2, 3], [4, 5, 6]].
      # Input activation at output layer i = x*W = [[9, 12, 15], [19, 26, 33]].
      # Logits output therefore becomes E*i, where E are the embeddings of output
      #  keys, i.e., E = [[1, 2, 3], [4, 5, 6], [7, 8, 9]].
      # Then the logits output becomes [[78, 186, 294], [170, 404, 638]]
      #
      # If we define the loss to be L = tf.reduced_sum(Logits), then
      # dL/dE = sum_by_key(i) = [[28, 38, 48], [28, 38, 48], [28, 38, 48]].
      # So the expected new embeddings become
      # E - 0.1 * dL/dE = [[-1.8, -1.8, -1.8], [1.2, 1.2, 1.2], [4.2, 4.2, 4.2]].
      weights = tf.Variable([[1, 2, 3], [4, 5, 6]], dtype=tf.float32)
      inputs = tf.constant([[1.0, 2.0], [3.0, 4.0]])
  
      with tf.GradientTape() as tape:
          logits, _, _, _, _ = cs_ops.compute_sampled_logits(
              [["key1", ""], ["key2", "key3"]],
              tf.matmul(inputs, weights),
              3,
              de_config,
              "emb",
              service_address=self._kbs_address,
          )
          loss = tf.reduce_sum(logits)
  
      # Applies the gradient descent.
      grads = tape.gradient(loss, weights)
  
      # The gradients updated by the knowledge bank.
      updated_embedding = de_ops.dynamic_embedding_lookup(
          ["key1", "key2", "key3"],
          de_config,
          "emb",
          service_address=self._kbs_address,
      )
      self.assertAllClose(
          updated_embedding,
          [[-1.8, -1.8, -1.8], [1.2, 1.2, 1.2], [4.2, 4.2, 4.2]],
      )
  
      # The gradients w.r.t. the weight W is calculated as
      # dL/dw = dL/di * di/dW = sum_by_dim(E) * x =
      # [12, 15, 18] * [[4, 4, 4], [6, 6, 6]] = [[48, 60, 72], [72, 90, 108]]
      self.assertAllClose(grads, [[48, 60, 72], [72, 90, 108]])
  
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 5 for calling the CandidateSamplingOpsTest obj's underlying member method code with example usage: obj.test_compute_sampled_logits_grad() without return types