Replace upstream PyTorch GRU module

[tengyifei]

Now that we have a GRU implemented with scan (#8777), we can figure out a way to patch or intercept the PyTorch upstream GRU module such that it just gets replaced with our own GRU logic. Then the user won't have to know and the upstream GRU just works.

[iwknow]

i can take a look at this issue if no one else is assigned.

[tengyifei]

Hi @iwknow , thanks for offering to take a look. Here are more details. If you take a look at #8777, it adds a GRU implementation that is interface-compatible with upstream PyTorch, but not weight compatible. Also it doesn't support batch_first and bidirectional.

My thoughts are:

• We'd need to make the experimental GRU weight key-compatible with PyTorch also. Otherwise existing checkpoints cannot be loaded.
• Maybe we need to teach the experimental GRU to fallback to upstream GRU somehow if the user passes batch_first or bidirectional, or support these in the scan-based GRU
• Once we make our GRU fully compatible with upstream, we could add a patching logic to https://github.com/pytorch/xla/blob/master/torch_xla/_patched_functions.py#L60

[iwknow]

this issue may be closed.

[tengyifei]

Thanks @iwknow !