# Path: mesh_tensorflow/transformer/memory_layers.py
# class ProductKeyValueMemory(transformer.TransformerLayer):
#   def __init__(self,
#                key_size: int = gin.REQUIRED,
#                n_keys: int = gin.REQUIRED,
#                n_heads: int = gin.REQUIRED,
#                knn: int = gin.REQUIRED):
#   def call(self, context, x: mtf.Tensor) -> mtf.Tensor:
#   def get_indices(self, keys: mtf.Tensor,
#                   query: mtf.Tensor) -> Tuple[mtf.Tensor, mtf.Tensor]:


import mesh_tensorflow as mtf
import mock
import numpy as np
import tensorflow.compat.v1 as tf
from mesh_tensorflow.transformer import memory_layers
            [[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)


# Path: mesh_tensorflow/transformer/memory_layers.py
# class ProductKeyValueMemory(transformer.TransformerLayer):
#   def __init__(self,
#                key_size: int = gin.REQUIRED,
#                n_keys: int = gin.REQUIRED,
#                n_heads: int = gin.REQUIRED,
#                knn: int = gin.REQUIRED):
#   def call(self, context, x: mtf.Tensor) -> mtf.Tensor:
#   def get_indices(self, keys: mtf.Tensor,
#                   query: mtf.Tensor) -> Tuple[mtf.Tensor, mtf.Tensor]:


# Path: mesh_tensorflow/transformer/memory_layers_test.py
import mesh_tensorflow as mtf
import mock
import numpy as np
import tensorflow.compat.v1 as tf
from mesh_tensorflow.transformer import memory_layers

            [[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)