scan-based GRU falls back to nn.GRU when bidirectional is true. (#8984)

Author: iwknow

test/test_gru.py

@@ -14,27 +14,31 @@ def setUp(self):
     torch.manual_seed(0)
     torch_xla.manual_seed(0)
 
-  def build_models(self,
-                   input_size,
-                   hidden_size,
-                   num_layers,
-                   bias,
-                   batch_first=False):
+  def build_models(
+      self,
+      input_size,
+      hidden_size,
+      num_layers,
+      bias,
+      batch_first=False,
+      bidirectional=False,
+  ):
     gru = nn.GRU(
         input_size,
         hidden_size,
         num_layers=num_layers,
         bias=bias,
         batch_first=batch_first,
         dropout=0.0,
-        bidirectional=False)
+        bidirectional=bidirectional)
     scan_gru = GRU(
         input_size,
         hidden_size,
         num_layers=num_layers,
         bias=bias,
         batch_first=batch_first,
-        dropout=0.0)
+        dropout=0.0,
+        bidirectional=bidirectional)
 
     # Copy parameters from the upstream GRU to our scan-based GRU.
     # This ensures that the scan-based GRU has the same parameters as the
@@ -91,6 +95,22 @@ def check_gradients(self,
             atol=atol,
             rtol=rtol)
 
+  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, (
+        "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, (
+        "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.
@@ -114,7 +134,10 @@ def test_scan_gru_and_upstream_gru_interchangeability(self):
 
     # Check that the methods of the GRU and scan-based GRU have the same signature.
     common_methods = nn_gru_names & scan_gru_names
+    exempt_methods = ['__new__']
     for method_name in common_methods:
+      if method_name in exempt_methods:
+        continue
       try:
         nn_gru_method = nn_gru_members[method_name]
         scan_gru_method = scan_gru_members[method_name]

torch_xla/experimental/gru.py

@@ -1,6 +1,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import logging
 from torch.nn.utils.rnn import PackedSequence
 from typing import overload
 
@@ -59,6 +60,15 @@ class GRU(nn.GRU):
 
   """
 
+  def __new__(cls, *args, **kwargs):
+    if ('bidirectional' in kwargs and kwargs['bidirectional'] == True):
+      logging.warning(
+          "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)
+    return super().__new__(cls)
+
   @overload
   def __init__(
       self,