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
Cheers