class_docstr: Memory network with product key-value structure.
docstr: Call the layer.
code: def call(self, context, x: mtf.Tensor) -> mtf.Tensor:
      """Call the layer."""
      # Initialize Memory Keys and Values
      n_key_dim = mtf.Dimension("n_keys", self.n_keys)
      n_value_dim = mtf.Dimension("n_values", self.n_values)
      key_dim = mtf.Dimension("key", self.key_size // 2)
      value_dim = x.shape.dims[-1]
      head_dim = mtf.Dimension("n_heads", self.n_heads)
      product_dim = mtf.Dimension("product_key", 2)
      keys = mtf.get_variable(
          context.mesh,
          name="keys",
          shape=mtf.Shape([head_dim, product_dim, n_key_dim, key_dim]),
          dtype=context.variable_dtype,
      )
      values = mtf.layers.embedding_weights(
          context.mesh,
          vocab_dim=n_value_dim,
          output_dim=value_dim,
          variable_dtype=context.variable_dtype,
          name="values",
      )
  
      # Compute query
      new_dims = [head_dim, product_dim, key_dim]
      reduce_dims = x.shape.dims[-1:]
      query = mtf.layers.dense(
          x,
          new_dims,
          reduced_dims=reduce_dims,
          activation=None,
          use_bias=True,
          variable_dtype=context.variable_dtype,
          name="query",
      )  # [b, l, h, 2, k]
  
      # Note: We use layer norm instead of batch norm to normalize queries.
      # The main advantage is that layer norm works well with the codebase
      # whereas the implementation of batch norm requires handling of tf ops.
      query = mtf.layers.layer_norm(query, query.shape.dims[-1])
  
      # Retrieve indices and scores
      scores, indices = self.get_indices(keys, query)  # [b, l, h, k]
      scores = mtf.softmax(scores, reduced_dim=scores.shape.dims[-1])
      top_values = mtf.gather(
          values, indices, n_value_dim
      )  # [b, l, h, k, v]
      out_values = mtf.einsum(
          [top_values, scores], reduced_dims=scores.shape.dims[-2:]
      )  # [b, l, v]
      return out_values
  
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 2 for calling the ProductKeyValueMemory obj's underlying member method code with example usage: obj.call(context, x) and returns: out_values