code: def test_get_indices(self):
      key_size = 2
      n_keys = 3
      product_size = 2
      head_size = 2
      batch = 2
      seq_len = 2
      knn = 2
  
      n_key_dim = mtf.Dimension("n_keys", n_keys)
      key_dim = mtf.Dimension("key", key_size // 2)
      seq_dim = mtf.Dimension("length", seq_len)
      batch_dim = mtf.Dimension("batch", batch)
      head_dim = mtf.Dimension("n_heads", head_size)
      product_dim = mtf.Dimension("product_key", product_size)
      knn_dim = mtf.Dimension("knn", knn)
  
      query_shape = mtf.Shape(
          [batch_dim, seq_dim, head_dim, product_dim, key_dim]
      )
      keys_shape = mtf.Shape(
          [head_dim, product_dim, n_key_dim, key_dim]
      )
      query = mtf.ones(self.mesh, query_shape)
  
      keys_vals = [
          [
              [[4], [1], [2]],
              [[2], [-1], [2]],
          ],
          [
              [[1], [2], [5]],
              [[6], [1], [4]],
          ],
      ]
      # h1:
      #   First scores:
      #   [4, 2]
      #   [2, 2]
      #   Cartesian added scores:
      #   [6, 6]
      #   Indices:
      #   [0, 2]    [0*n_k + 0, 0*n_k + 2]
      # h2:
      #   First scores:
      #   [5, 2]
      #   [6, 4]
      #   Cartesian added scores:
      #   [11, 9]
      #   Indices:
      #   [6, 8]   [2*n_k+0, 2*n_k+2]
      expected_scores = np.broadcast_to(
          np.array([[6, 6], [11, 9]]), [batch, seq_len, head_size, knn]
      )
      expected_indices = np.broadcast_to(
          np.array([[0, 2], [6, 8]]), [batch, seq_len, head_size, knn]
      )
  
      keys = mtf.constant(self.mesh, keys_vals, keys_shape)
  
      pkm = memory_layers.ProductKeyValueMemory(
          key_size, n_keys, head_size, knn
      )
      mtf_scores, mtf_indices = pkm.get_indices(keys, query)
  
      # Shapes.
      expected_shape = mtf.Shape(
          [batch_dim, seq_dim, head_dim, knn_dim]
      )
      self.assertEqual(expected_shape, mtf_scores.shape)
      self.assertEqual(expected_shape, mtf_indices.shape)
  
      # Values
      lowering_s, scores = self._export_to_tf_tensor(mtf_scores)
      lowering_i, indices = self._export_to_tf_tensor(mtf_indices)
      self.evaluate(tf.global_variables_initializer())
      self.evaluate(lowering_s.copy_masters_to_slices())
      self.evaluate(lowering_i.copy_masters_to_slices())
      scores, indices = self.evaluate([scores, indices])
  
      self.assertAllEqual(expected_scores, scores)
      self.assertAllEqual(expected_indices, indices)
  
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 4 for calling the FlatKeyValueMemoryTest obj's underlying member method code with example usage: obj.test_get_indices() without return types