Replace upstream GRU implementation with scan-based GRU (#9010)

Author: iwknow

test/test_gru.py

@@ -2,7 +2,7 @@
 import torch.nn as nn
 import torch_xla
 import inspect
-from torch_xla.experimental.gru import GRU
+from torch_xla.experimental.gru import GRU as ScanGRU
 
 from absl.testing import absltest, parameterized
 
@@ -23,15 +23,15 @@ def build_models(
       batch_first=False,
       bidirectional=False,
   ):
-    gru = nn.GRU(
+    gru = nn.GRU._orig(
         input_size,
         hidden_size,
         num_layers=num_layers,
         bias=bias,
         batch_first=batch_first,
         dropout=0.0,
         bidirectional=bidirectional)
-    scan_gru = GRU(
+    scan_gru = nn.GRU(
         input_size,
         hidden_size,
         num_layers=num_layers,
@@ -95,28 +95,45 @@ def check_gradients(self,
             atol=atol,
             rtol=rtol)
 
+  def test_patch_happened(self):
+    """
+    Ensures that the GRU class is patched correctly. The patch should happen in _patched_functions.py before
+    this test is run.
+    """
+    # Check if the GRU class is patched.
+    assert type(nn.GRU) is type(ScanGRU), (
+        "GRU class should be patched. "
+        "Check if the patching code is executed before this test.")
+    assert hasattr(
+        nn.GRU,
+        '_orig'), ("GRU class should be patched. "
+                   "Check if the patching code is executed before this test.")
+    assert nn.GRU._orig is not None, (
+        "GRU class should have the original GRU class as _orig. "
+        "Check if the patching code is executed before this test.")
+
   def test_scan_gru_fallback_to_upstream_gru(self):
     """
     Ensures that the scan-based GRU falls back to the upstream GRU when
     unsupported parameters are set.
     """
     input_size, hidden_size, num_layers = 16, 32, 2
     _, scan_gru = self.build_models(input_size, hidden_size, num_layers, True)
-    assert type(scan_gru) is GRU, (
+    assert type(scan_gru) is nn.GRU, (
         "Scan-based GRU should create scan-based GRU when *no* unsupported parameters are set."
     )
     _, scan_gru = self.build_models(
         input_size, hidden_size, num_layers, True, bidirectional=True)
-    assert type(scan_gru) is nn.GRU, (
+    assert type(scan_gru) is nn.GRU._orig, (
         "Scan-based GRU should fall back to upstream GRU when `bidirectional` is set to True."
     )
 
   def test_scan_gru_and_upstream_gru_interchangeability(self):
     """
     Ensures that the scan-based GRU and upstream GRU are interchangeable.
     """
-    nn_gru = nn.GRU
-    scan_gru = GRU
+    nn_gru = nn.GRU._orig
+    scan_gru = nn.GRU
     nn_gru_members = dict(inspect.getmembers(nn_gru, inspect.isroutine))
     scan_gru_members = dict(inspect.getmembers(scan_gru, inspect.isroutine))
 

torch_xla/_patched_functions.py

@@ -3,11 +3,17 @@
 import torch.nn as nn
 from torch import inf
 from typing import Iterable, Union, Optional
+from torch_xla.experimental.gru import GRU as ScanGRU
 
 _tensor_or_tensors = Union[torch.Tensor, Iterable[torch.Tensor]]
 
 
-def _patch(fn, newfn):
+def _pathch_module(m, new_m):
+  new_m._orig = m
+  return new_m
+
+
+def _patch_fn(fn, newfn):
   xfingerprint = inspect.signature(fn)
   fingerprint = inspect.signature(newfn)
   if xfingerprint != fingerprint:
@@ -58,4 +64,6 @@ def clip_grad_norm_(parameters: _tensor_or_tensors,
 
 
 def _apply_patches():
-  nn.utils.clip_grad_norm_ = _patch(nn.utils.clip_grad_norm_, clip_grad_norm_)
+  nn.utils.clip_grad_norm_ = _patch_fn(nn.utils.clip_grad_norm_,
+                                       clip_grad_norm_)
+  nn.GRU = _pathch_module(nn.GRU, ScanGRU)

torch_xla/experimental/gru.py

@@ -66,7 +66,15 @@ def __new__(cls, *args, **kwargs):
           "Scan-based GRU only supports unidirectional GRU. (bidirectional = False) "
           "Scan-based GRU falls back to the default nn.GRU implementation instead."
       )
-      return nn.GRU(*args, **kwargs)
+      if nn.GRU._orig is None:
+        # If nn.GRU._orig is None, it means that the original GRU has not been
+        # patched yet for some reason. The patching should happen in _patched_functions.py.
+        # So we need to call the original GRU constructor here.
+        return nn.GRU(*args, **kwargs)
+      else:
+        # If nn.GRU._orig is not None, it means that the original GRU has been
+        # patched already. So we need to call the patched GRU constructor here.
+        return nn.GRU._orig(*args, **kwargs)
     return super().__new__(cls)
 
   @overload