code: def testWrite(self):
      batch_size = 2
      key_depth = 3
      val_depth = 5
      memory_size = 4
      window_size = 6
      x_depth = 10
      memory = transformer_memory.TransformerMemory(
          batch_size, key_depth, val_depth, memory_size
      )
      x = tf.random_uniform(
          [batch_size, window_size, x_depth], minval=1.0
      )
      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)
      with tf.control_dependencies([update_op]):
          logits, _ = memory.read(x)
          write_op = memory.write(x, logits)
      mem_vals, mem_logits = memory.get()
      with self.test_session() as session:
          session.run(tf.global_variables_initializer())
          session.run(write_op)
          updated_vals, updated_logits = session.run(
              [mem_vals, mem_logits]
          )
      self.assertAllEqual(
          [batch_size, memory_size, val_depth], updated_vals.shape
      )
      self.assertAllEqual(
          [batch_size, memory_size], updated_logits.shape
      )
  
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 2 for calling the TransformerMemoryTest obj's underlying member method code with example usage: obj.testWrite() without return types