Merge pull request #2610 from chiamp:test_returned_variable_warning

PiperOrigin-RevId: 488461724
google · Nov 14, 2022 · 1a96148 · 1a96148
2 parents 478cffe + f791354
commit 1a96148
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/tests/linen/linen_transforms_test.py b/tests/linen/linen_transforms_test.py
@@ -1026,6 +1026,19 @@ def _helper(self):
     with self.assertRaises(errors.TransformedMethodReturnValueError):
       b.apply({}, jnp.ones(2))
 
+  def test_returned_variable_warning(self):
+    class Bar(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        f = self._helper()
+        return f(x)
+      @nn.jit
+      def _helper(self):
+        return nn.Variable(None, None, None)
+    b = Bar()
+    with self.assertRaises(errors.TransformedMethodReturnValueError):
+      b.apply({}, jnp.ones(2))
+
   def test_nowrap(self):
     class Bar(nn.Module):
       @nn.compact