class_docstr: Memory network with product key-value structure.
docstr: Generate score and indices for the query.
code: def get_indices(
      self, keys: mtf.Tensor, query: mtf.Tensor
  ) -> Tuple[mtf.Tensor, mtf.Tensor]:
      """Generate score and indices for the query."""
      score_shape = mtf.Shape(
          query.shape.dims[:-1] + keys.shape.dims[2:3]
      )
      scores = mtf.einsum(
          [query, keys], output_shape=score_shape
      )  # [b, l, h, 2, n_keys]
      knn_dim = mtf.Dimension("knn", self.knn)
      scores, indices = mtf.top_k(
          scores, score_shape.dims[-1], knn_dim
      )  # [b, l, h, 2, knn]
  
      # Computes the top cartesian products and their indices
      knn_square_dim = mtf.Dimension("knn_square_dim", self.knn**2)
      scores1, scores2 = mtf.unstack(scores, scores.shape.dims[-2])
      scores2 = mtf.rename_dimension(scores2, "knn", "knn2")
      out_shape = mtf.Shape(
          scores1.shape.dims + scores2.shape.dims[-1:]
      )
      all_scores = mtf.add(scores1, scores2, output_shape=out_shape)
      all_scores = mtf.replace_dimensions(
          all_scores, out_shape[-2:], knn_square_dim
      )
  
      indices1, indices2 = mtf.unstack(indices, indices.shape.dims[-2])
      indices1 = mtf.multiply(indices1, self.n_keys)
      indices2 = mtf.rename_dimension(indices2, "knn", "knn2")
      all_indices = mtf.add(indices1, indices2, output_shape=out_shape)
      all_indices = mtf.replace_dimensions(
          all_indices, out_shape[-2:], knn_square_dim
      )
  
      scores, best_indices = mtf.top_k(
          all_scores, all_scores.shape.dims[-1], knn_dim
      )
      return scores, mtf.gather(
          all_indices, best_indices, knn_square_dim
      )
  
file_path: mesh/mesh_tensorflow/transformer/memory_layers.py
imports: from typing import Tuple
  import gin
  import mesh_tensorflow
  from mesh_tensorflow.transformer import transformer
label_desc: Use image node_id 3 for calling the ProductKeyValueMemory obj's underlying member method code with example usage: obj.get_indices(keys, query) and returns: scores,  mtf