class_docstr: Memory network with product key-value structure.
docstr: Creates a ProductKeyValueMemory layer.
code: def __init__(
      self,
      key_size: int = gin.REQUIRED,
      n_keys: int = gin.REQUIRED,
      n_heads: int = gin.REQUIRED,
      knn: int = gin.REQUIRED,
  ):
      """Creates a ProductKeyValueMemory layer."""
      self.key_size = key_size
      self.n_keys = n_keys
      self.n_values = n_keys**2
      self.n_heads = n_heads
      self.knn = knn
  
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 1 to create a new ProductKeyValueMemory object from inherited base classes: transformer with example: obj = ProductKeyValueMemory(key_size, n_keys, n_heads, knn)