code: def testReset(self):
      batch_size = 2
      key_depth = 3
      val_depth = 5
      memory_size = 4
      memory = transformer_memory.TransformerMemory(
          batch_size, key_depth, val_depth, memory_size
      )
      vals = tf.random_uniform(
          [batch_size, memory_size, val_depth], minval=1.0
      )
      logits = tf.random_uniform([batch_size, memory_size], minval=1.0)
      update_op = memory.set(vals, logits)
      reset_op = memory.reset([1])
      mem_vals, mem_logits = memory.get()
      assert_op1 = tf.assert_equal(mem_vals[0], vals[0])
      assert_op2 = tf.assert_equal(mem_logits[0], logits[0])
      with tf.control_dependencies([assert_op1, assert_op2]):
          all_zero1 = tf.reduce_sum(tf.abs(mem_vals[1]))
          all_zero2 = tf.reduce_sum(tf.abs(mem_logits[1]))
      with self.test_session() as session:
          session.run(tf.global_variables_initializer())
          session.run(update_op)
          session.run(reset_op)
          zero1, zero2 = session.run([all_zero1, all_zero2])
      self.assertAllEqual(0, zero1)
      self.assertAllEqual(0, zero2)
  
file_path: tensor2tensor/tensor2tensor/layers/transformer_memory_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  from tensor2tensor.layers import transformer_memory
  import tensorflow.compat.v1
label_desc: Use image node_id 3 for calling the TransformerMemoryTest obj's underlying member method code with example usage: obj.testReset() without return types