code: def test_call_shape(self):
      key_size = 5
      value_size = 10
      n_keys = 6
      n_heads = 2
      knn = 3
  
      seq_len = 4
      batch = 5
  
      model_dim = mtf.Dimension("model", value_size)
      seq_dim = mtf.Dimension("length", seq_len)
      batch_dim = mtf.Dimension("batch", batch)
  
      def initialize(shape, dtype):
          return tf.reshape(
              1 + tf.range(np.prod(shape), dtype=dtype), shape
          )
  
      self.initializer_mock.side_effect = initialize
  
      kv_memory = memory_layers.ProductKeyValueMemory(
          key_size, n_keys, n_heads, knn
      )
  
      mtf_x = mtf.ones(
          self.mesh, mtf.Shape([batch_dim, seq_dim, model_dim])
      )
      context = mock.MagicMock()
      context.mesh = self.mesh
      context.variable_dtype = tf.float32
  
      out_tensor = kv_memory.call(context, mtf_x)
  
      # Dimensions should be untouched
      self.assertEqual(mtf_x.shape, out_tensor.shape)
  
file_path: mesh/mesh_tensorflow/transformer/memory_layers_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import mesh_tensorflow
  from mesh_tensorflow.transformer import memory_layers
  import mock
  import numpy
  import tensorflow.compat.v1
label_desc: Use image node_id 3 for calling the FlatKeyValueMemoryTest obj's underlying member method code with example usage: obj.test_call_shape() and returns: tf