GRU/RNN state in graph mode Vs eager execution mode - tensorflow

Question

I have same piece of code written first in eager execution mode and then in graph mode. Now, I am not quite able to figure out why the GRU state is not retained in the graph mode while it's working fine in eager mode.
Here's the eager mode code:
import tensorflow as tf
import xxhash
import numpy as np
tf.enable_eager_execution()
rnn_units = 1024
def hash_code(arr):
return xxhash.xxh64(arr).hexdigest()
model = tf.keras.Sequential([tf.keras.layers.GRU(rnn_units,
return_sequences=True,
stateful=True,
recurrent_initializer='glorot_uniform', batch_input_shape=[1, None, 256])])
lstm_wt = np.load('lstm_wt.npy', allow_pickle=True) # fixed weights for comparison
lstm_re_wt = np.load('lstm_re_wt.npy', allow_pickle=True)
lstm_bias = np.load('lstm_bias.npy', allow_pickle=True)
model.layers[0].set_weights([lstm_wt, lstm_re_wt, lstm_bias])
op_embed = np.load('op_embed.npy', allow_pickle=True) # fixed input
op_lstm = model(op_embed)
print(hash_code(op_lstm.numpy()))
op_lstm = model(op_embed)
print(hash_code(op_lstm.numpy()))
model.layers[0].reset_states() # now reset the state, you'll get back the initial output.
op_lstm = model(op_embed)
print(hash_code(op_lstm.numpy()))
Output of this code:
d092fdb4739588a3
cdfdf8b8e292c6e8
d092fdb4739588a3
Now, the graph model code:
import tensorflow as tf
import xxhas
import numpy as np
# checking lstm
op_embed = np.load('op_embed.npy', allow_pickle=True)
# load op_embed, lstm weights
lstm_wt = np.load('lstm_wt.npy', allow_pickle=True)
lstm_re_wt = np.load('lstm_re_wt.npy', allow_pickle=True)
lstm_bias = np.load('lstm_bias.npy', allow_pickle=True)
rnn_units = 1024
layers = tf.keras.layers.GRU(rnn_units,
return_sequences=True,
stateful=True,
recurrent_initializer='glorot_uniform')
x_placeholder = tf.placeholder(shape=op_embed.shape, dtype=tf.float32)
op_lstm = layers(x_placeholder)
sess = tf.Session()
sess.run(tf.global_variables_initializer())
layers.set_weights([lstm_wt, lstm_re_wt, lstm_bias])
tf.assign(layers.weights[0],lstm_wt ).eval(sess)
tf.assign(layers.weights[1], lstm_re_wt).eval(sess)
tf.assign(layers.weights[2], lstm_bias).eval(sess)
print('keras op hash',xxhash.xxh64(sess.run(op_lstm, feed_dict={x_placeholder:op_embed})).hexdigest())
print('keras op hash',xxhash.xxh64(sess.run(op_lstm, feed_dict={x_placeholder:op_embed})).hexdigest())
output:
keras op hash d092fdb4739588a3
keras op hash d092fdb4739588a3
Any insights on how to fix this ambiguity and retain states in case of graph mode?
There's a similar question asked before but unanswered. Statefulness in eager mode vs non-eager mode

Answer 1

Specifying the Solution here (Answer Section) even though it is present in the Link provided in the Question, for the Benefit of the Community.
The Recurrent Neural Network (RNN or GRU or LSTM) loses its state while executing it in Non-Eager-Mode/Graph-Mode by default.
If we want to retain the state, we need to pass the Initial State during RNN call, as shown below:
current_state = np.zeros((1,1))
state_placeholder = tf.placeholder(tf.float32, shape=[1, 1])
output, state = rnn(x, initial_state=state_placeholder)
Then, while executing the Output, we need to pass the State as well, in addition to the Input for feed_dict.
So, the code,
print('keras op hash',xxhash.xxh64(sess.run(op_lstm, feed_dict={x_placeholder:op_embed})).hexdigest())
can be replaced with
for _ in range(No_Of_TimeSteps):
op_val, state_val = sess.run([op_lstm, state], feed_dict={x_placeholder:op_embed})).hexdigest(),
state_placeholder: current_state.astype(np.float32)})
current_state = state_val
print('keras op hash',xxhash.xxh64(op_val))
Hope this helps. Happy Learning!