(Vanilla) RNN to FTRNN Implementation

Hey

I want to change (Vanilla) RNN model by applying Tensor Factorization method on it. The equations are as:

(Vanilla) RNN:
h(t) = tanh( W[hx].x(t) + U[hh].h(t) + b(h))

FTRNN:
h(t) = tanh( W[hfx] diag (W[fxi].I) W[fxx].x(t) + U[hfh] diag (U[fhi].I) U[fhh].h(t-1) + b(h))

where as, W[hfx] has dimensions nh x nf ; W[fxi] has dimensions nf x |I| ; W[fxx] has dimensions nf x nx and similarly for hidden to hidden layer

fx & fh are tensor factors for input and hidden layer respectively for input to hidden layer (as represented by W) and hidden to hidden layer ( as represented by U).

Hyper parameters are nh = hidden size, nfx = factor size(input to hidden layer), nfh = factor size (hidden-hidden layer)
Moreover,x = input features (given to each time step) to model and I = constant/scaler (it is also given to model)

I am beginner in pytorch. Would some one please help me or have any suggestion to implement FTRNN in pytorch or should I have to change (Source code for torch.nn.modules.rnn) ?

I shall be very thankful to you

look at Implementation of Multiplicative LSTM for some pointers on implementing custom rnns.