Simple GRU initialization not working

I tried to pass just 1 word input to a GRU layer. Below is what I did -


But I got error “Number of weight elements (6) does not equal the number of sublayers (3) in: GRU_5.

Also, running model.sublayers gave below



What am I doing wrong? What do these sublayer values mean? Where will the model reflect n_units = 5 that I gave during initialization ?

Hi @Mayank11

There is a lot to explain, you better read the docs and check the code.

I think for the start what you want is:

model = tl.GRU(n_units=5)
x = np.array([[[1,2,3,4,5]]])


This is very different from what was taught in the lab. The lab taught GRU as below -

Then 1 forward pass for 1st time step was done as follows -


And then the lab told to implement a scan function that does above for all elements in the input.

I guess out-of-the-box GRU in trax library behaves differently…

For the example code shared by you, following is the output I am getting. Can you help me interpret what below shapes mean -



The lab implementation is actually very accurate, it’s just not exactly the same as in the trax library. Depending on what are you trying to achieve, you can use the lab implementation to your original quest and everything should work and you wouldn’t need to read the code.

But if you want to understand the trax implementation, there is just simply too much to explain in a single post. I would encourage you to read the code line by line.

Here you have:

  • weights[0][0], shape (10, 10) - Update and reset gates weights (combined)
  • weights[0][1], shape (10, ) - Update and reset gates bias
  • weights[0][2], shape (10, 5) - Candidate weights
  • weights[0][3], shape (5, ) - Candidate bias