Batch_RWV_Generation.py

import tensorflow as tf
import numpy as np

@tf.function
def ReadVector(M_t, w_t):
    '''
    Computes the Read Vector of EACH HEAD for the entire Batch at once.
    
    M_t: Memory of size (Batch_size,N,M) at time t.
    w_t: (Batch_size,N), Weighting generated by READ HEAD at time t for reading memory location
    
    RETURNS:
    
    r_t: (Batch_size,M) The Read Vector
    '''
    batch_size,N,M = M_t.shape[0], M_t.shape[1], M_t.shape[2]
    
    #tol = 0.01
    #assert (np.sum(w_t) >= 1.0 - tol) & (np.sum(w_t) <= 1.0 + tol)
    #print(batch_size,N,M)
    #print('w_t shape: ',w_t.shape)
    r_t = tf.reshape(  tf.matmul(tf.reshape(w_t,(batch_size,1,N)),M_t), (batch_size,M)   )
    
    #tf.print('ReadVector: ',tf.math.reduce_any(tf.math.is_nan( r_t )) )
    #assert r_t.shape == (batch_size,M)
    
    return r_t

@tf.function
def WriteOnMemory(M_prev, w_t, e_t, a_t):
    
    '''
    Computes the updated Memory Matrix for the each example in Batch at once
    
    M_prev: Memory Matrix at the previous time step of size (Batch_size,N,M)
    w_t: (Batch_size,N), Weighting generated by WRITE HEAD at time t for Writing to the memory locations.
    e_t: (Batch_size,M), Erase vector generated by WRITE HEAD.
    a_t: (Batch_size,M), Add vector generated by WRITE HEAD.
    
    RETURNS:
    
    M_t: (Batch_size,N,M), New Memory Matrix after Erasing/Adding/Combination of new instances.
    '''
    
    (batch_size,N,M) = M_prev.shape
    
    M_hat_t = tf.multiply( M_prev, 1-tf.multiply(tf.reshape(w_t,(batch_size,N,1)),tf.reshape(e_t,(batch_size,1,M))))
    #^Of shape [batch_size,N,M]
    
    #tf.print('WriteOnMemory M_hat_t',tf.math.reduce_any(tf.math.is_nan( M_hat_t )) )

    
    #assert M_hat_t.shape == M_prev.shape
    
    M_t = M_hat_t + tf.multiply(tf.reshape(w_t,(batch_size,N,1)),tf.reshape(a_t,(batch_size,1,M)))
    
    #tf.print('WriteOnMemory',tf.math.reduce_any(tf.math.is_nan( M_t )) )

    
    #assert M_t.shape == M_prev.shape
    
    return M_t