# Path: src/modules/recurring.py
# class Recurring(Module):
#     """
#     Recurring is a type of module that cache intermediate
#     activations in a stack during unrolling of recurrent layers
#     Its backward() is expected to be called several times during
#     Back Propagation Through Time, either full or truncated.
#     """
#     def __init__(self, *args, **kwargs):
#         self._stack = list()
#         self._setup(*args, **kwargs)
# 
#     def _setup(*args, **kwargs):
#         pass
# 
#     def _push(self, *objs):
#         objs = list(objs)
#         if len(objs) == 1:
#             objs = objs[0]
#         self._stack.append(objs)
#     
#     def _pop(self):
#         objs = self._stack[-1]
#         del self._stack[-1]
#         return objs
#     
#     def flush(self):
#         self._stack = list()
#         self._flush()
#     
#     def _flush(self): pass
# 
# class gate(Recurring):
#     """ 
#     Gates are special case of Recurring modules
#     It starts with a linear transformation (matmul)
#     And ends with a (typically) nonlinear activation
#     """
#     def _setup(self, server, w_shape, bias = None, 
#                act_class = sigmoid, transfer = None):
#         if transfer is None:
#             b_shape = (w_shape[-1],)
#             w_init = xavier(w_shape)
#             if bias is not None:
#                 b_init = np.ones(b_shape) * bias
#             elif bias is None:
#                 b_init = guass(0., 1e-1, b_shape)
#         else: w_init, b_init = transfer
#         
#         self._act_class = act_class
#         self._w = server.issue_var_slot(w_init, True)
#         self._b = server.issue_var_slot(b_init, True)
#     
#     def forward(self, x):
#         linear = x.dot(self._w.val) + self._b.val
#         act = self._act_class(None, None)
#         self._push(x, act)
#         return act.forward(linear)
#         
#     def backward(self, grad):
#         x, act = self._pop()
#         linear_grad = act.backward(grad)
# 
#         self._b.set_grad(linear_grad.sum(0))
#         self._w.set_grad(x.T.dot(linear_grad))
#         return linear_grad.dot(self._w.val.T)
# 
# Path: src/modules/activations.py
# class softplus(Activate):
#     def transform(self, x):
#         self.activation = np.log(1 + np.exp(-x))
#     
#     def backward(self, grad):
#         a = self.activation
#         p = 1 - np.exp(-a)
#         return np.multiply(grad, p)
# 
# class tanh(Activate):
#     def transform(self, x):
#         self.activation = np.tanh(x)
#     
#     def backward(self, grad):
#         a = self.activation
#         p = 1 - np.multiply(a, a)
#         return np.multiply(grad, p)
# 
# class sigmoid(Activate):
#     def transform(self, x):
#         self.activation = 1. / (1. + np.exp(-x))
# 
#     def backward(self, grad):
#         a = self.activation
#         p = np.multiply(a, 1. - a)
#         return np.multiply(grad, p)
# 
# class softmax(Activate):
#     def transform(self, x):
#         row_max = x.max(1, keepdims = True)
#         e_x = np.exp(x - row_max)
#         e_sum = e_x.sum(1, keepdims = True)
#         self.activation = np.divide(e_x, e_sum)
# 
#     def backward(self, grad):
#         a = self.activation
#         m = np.multiply(grad, a)
#         g = grad - m.sum(1, keepdims = True)
#         return np.multiply(g, a)
# 
# class relu(Activate):
#     def transform(self, x):
#         self.activation = x * (x > 0.)
# 
#     def backward(self, grad):
#         p = self.activation > 0.
#         return np.multiply(grad, p)


from .recurring import Recurring, gate
from .activations import softplus, tanh, sigmoid, softmax, relu
from .mechanics import * 

class ntm_attend(Recurring):
    def _setup(self, server, lstm_size, vec_size, shift):
        self._gates = dict({
            'k': gate(server, (lstm_size, vec_size), None, sigmoid), # key gate
            'b': gate(server, (lstm_size, 1), None, softplus), # strength gate
            'i': gate(server, (lstm_size, 1), None, sigmoid), # interpolate gate
            'r': gate(server, (lstm_size, 

# Path: src/modules/recurring.py
# class Recurring(Module):
#     """
#     Recurring is a type of module that cache intermediate
#     activations in a stack during unrolling of recurrent layers
#     Its backward() is expected to be called several times during
#     Back Propagation Through Time, either full or truncated.
#     """
#     def __init__(self, *args, **kwargs):
#         self._stack = list()
#         self._setup(*args, **kwargs)
# 
#     def _setup(*args, **kwargs):
#         pass
# 
#     def _push(self, *objs):
#         objs = list(objs)
#         if len(objs) == 1:
#             objs = objs[0]
#         self._stack.append(objs)
#     
#     def _pop(self):
#         objs = self._stack[-1]
#         del self._stack[-1]
#         return objs
#     
#     def flush(self):
#         self._stack = list()
#         self._flush()
#     
#     def _flush(self): pass
# 
# class gate(Recurring):
#     """ 
#     Gates are special case of Recurring modules
#     It starts with a linear transformation (matmul)
#     And ends with a (typically) nonlinear activation
#     """
#     def _setup(self, server, w_shape, bias = None, 
#                act_class = sigmoid, transfer = None):
#         if transfer is None:
#             b_shape = (w_shape[-1],)
#             w_init = xavier(w_shape)
#             if bias is not None:
#                 b_init = np.ones(b_shape) * bias
#             elif bias is None:
#                 b_init = guass(0., 1e-1, b_shape)
#         else: w_init, b_init = transfer
#         
#         self._act_class = act_class
#         self._w = server.issue_var_slot(w_init, True)
#         self._b = server.issue_var_slot(b_init, True)
#     
#     def forward(self, x):
#         linear = x.dot(self._w.val) + self._b.val
#         act = self._act_class(None, None)
#         self._push(x, act)
#         return act.forward(linear)
#         
#     def backward(self, grad):
#         x, act = self._pop()
#         linear_grad = act.backward(grad)
# 
#         self._b.set_grad(linear_grad.sum(0))
#         self._w.set_grad(x.T.dot(linear_grad))
#         return linear_grad.dot(self._w.val.T)
# 
# Path: src/modules/activations.py
# class softplus(Activate):
#     def transform(self, x):
#         self.activation = np.log(1 + np.exp(-x))
#     
#     def backward(self, grad):
#         a = self.activation
#         p = 1 - np.exp(-a)
#         return np.multiply(grad, p)
# 
# class tanh(Activate):
#     def transform(self, x):
#         self.activation = np.tanh(x)
#     
#     def backward(self, grad):
#         a = self.activation
#         p = 1 - np.multiply(a, a)
#         return np.multiply(grad, p)
# 
# class sigmoid(Activate):
#     def transform(self, x):
#         self.activation = 1. / (1. + np.exp(-x))
# 
#     def backward(self, grad):
#         a = self.activation
#         p = np.multiply(a, 1. - a)
#         return np.multiply(grad, p)
# 
# class softmax(Activate):
#     def transform(self, x):
#         row_max = x.max(1, keepdims = True)
#         e_x = np.exp(x - row_max)
#         e_sum = e_x.sum(1, keepdims = True)
#         self.activation = np.divide(e_x, e_sum)
# 
#     def backward(self, grad):
#         a = self.activation
#         m = np.multiply(grad, a)
#         g = grad - m.sum(1, keepdims = True)
#         return np.multiply(g, a)
# 
# class relu(Activate):
#     def transform(self, x):
#         self.activation = x * (x > 0.)
# 
#     def backward(self, grad):
#         p = self.activation > 0.
#         return np.multiply(grad, p)


# Path: src/modules/attention.py
from .recurring import Recurring, gate
from .activations import softplus, tanh, sigmoid, softmax, relu
from .mechanics import * 


class ntm_attend(Recurring):
    def _setup(self, server, lstm_size, vec_size, shift):
        self._gates = dict({
            'k': gate(server, (lstm_size, vec_size), None, sigmoid), # key gate
            'b': gate(server, (lstm_size, 1), None, softplus), # strength gate
            'i': gate(server, (lstm_size, 1), None, sigmoid), # interpolate gate
            'r': gate(server, (lstm_size, 

                              2 * shift + 1), None, softmax), # rotate gate