How can I implement a custom RNN (specifically an ESN) in Tensorflow?

Question

I am trying to define my own RNNCell (Echo State Network) in Tensorflow, according to below definition.

x(t + 1) = tanh(Win*u(t) + W*x(t) + Wfb*y(t))

y(t) = Wout*z(t)

z(t) = [x(t), u(t)]

x is state, u is input, y is output. Win, W, and Wfb are not trainable. All weights are randomly initialized, but W is modified like this: "Set a certain percentage of elements of W to 0, scale W to keep its spectral radius below 1.0

I have this code to generate the equation.

x = tf.Variable(tf.reshape(tf.zeros([N]), [-1, N]), trainable=False, name="state_vector") W = tf.Variable(tf.random_normal([N, N], 0.0, 0.05), trainable=False) # TODO: setup W according to the ESN paper W_x = tf.matmul(x, W)  u = tf.placeholder("float", [None, K], name="input_vector") W_in = tf.Variable(tf.random_normal([K, N], 0.0, 0.05), trainable=False) W_in_u = tf.matmul(u, W_in)  z = tf.concat(1, [x, u]) W_out = tf.Variable(tf.random_normal([K + N, L], 0.0, 0.05)) y = tf.matmul(z, W_out) W_fb = tf.Variable(tf.random_normal([L, N], 0.0, 0.05), trainable=False) W_fb_y = tf.matmul(y, W_fb)  x_next = tf.tanh(W_in_u + W_x + W_fb_y)  y_ = tf.placeholder("float", [None, L], name="train_output") 

My problem is two-fold. First I don't know how to implement this as a superclass of RNNCell. Second I don't know how to generate a W tensor according to above specification.

Any help about any of these question is greatly appreciated. Maybe I can figure out a way to prepare W, but I sure as hell don't understand how to implement my own RNN as a superclass of RNNCell.

Answer 1

To give a quick summary:

Look in the TensorFlow source code under python/ops/rnn_cell.py too see how to subclass RNNCell. It's usually like this:

class MyRNNCell(RNNCell):   def __init__(...):    @property   def output_size(self):   ...    @property   def state_size(self):   ...    def __call__(self, input_, state, name=None):      ... your per-step iteration here ... 

Answer 2

I may arrive a bit late but for anybody trying to do the same, tensorflow-addons added an implementation of ESNCell and ESN layer.