Fix GRU to match pytorch (#2701).

[nwhitehead]

Pull Request Template

Checklist

• Confirmed that run-checks all script has been executed.
• Made sure the book is up to date with changes in this PR.

Related Issues/PRs

This addresses issue #2701

Changes

• Update GRU implementation of "new" gate to match pytorch implementation. This can change numerical output in some cases.

• Add GRU unit test with sequence length > 1.

• Fix GRU input state dimensions and hidden state handling. This is an API change since the dimensions of the optional hidden state input are being corrected to the right sizes.

These changes do affect numerical results and change the API slightly. I think just updating to the correct API dimensions seems like the best thing since the previous implementation was incorrect, not just different than pytorch.

Testing

These changes were tested with a small unit test. For this test the correct values were computed manually using the equations for GRU.

I tested these changes against PyTorch. The weights and biases from PyTorch were saved then split into sections using a custom script (to split apart the weights for each gate). Input and output tensors were separately saved and then loaded into a test rust program. Everything was randomly initialized. With this PR the results from burn and torch were almost identical (within 6 decimal digits). I tried input sizes of 1, 2, and 8. I tried hidden sizes of 1, 2, and 8. I tried sequence lengths of 1, 2, and 3.

[codecov]

Codecov Report

Attention: Patch coverage is 91.57895% with 8 lines in your changes missing coverage. Please review.

Project coverage is 83.20%. Comparing base (f630b3b) to head (cc4766f).
Report is 3 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/burn-core/src/nn/rnn/gru.rs	91.57%	8 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2704   +/-   ##
=======================================
  Coverage   83.20%   83.20%           
=======================================
  Files         819      819           
  Lines      106814   106866   +52     
=======================================
+ Hits        88870    88916   +46     
- Misses      17944    17950    +6     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[laggui]

Thanks for bringing this up and tackling it yourself 🙏

Fun fact, the reset gate changes don't originate from pytorch 😄

The implementation we had is based on the latest v3 revisions (published at EMNLP) and has reset gate applied to hidden state before matrix multiplication.

The changes in your PR are based on the original v1 and applies the reset gate after.

Curiously, pytorch notes efficiency for their differing implementation (without much explanation). If you're interested, check out this awesome explanation behind the motivation to move the reset gate.

Your implementation LGTM! But what do you think about supporting both via a config? And we could provide the references I just linked in the doc.

[laggui]

Thank you 🙏

I just had to fix the clippy doc issue and covered both versions in the tests.