Handle a special case of lambda usage properly for tf.function freevar detection

PiperOrigin-RevId: 481253023

tensorflow/core/function/capture/free_vars_detect.py

@@ -124,6 +124,10 @@ def _search_callable_free_vars(fn):
   except ValueError:
     # When source code unavailable, return empty result
     return []
+  except NotImplementedError:
+    # Autograph cannot handle multiple lambda functions with same line number
+    # and args name.
+    return []
 
   scope = anno.getanno(node, anno.Static.SCOPE)
   free_vars_all = list(scope.free_vars)

tensorflow/core/function/capture/free_vars_detect_test.py

@@ -124,6 +124,12 @@ class FreeVarDetectionTest(parameterized.TestCase):
     free_vars = get_var_name(func_map["f"])
     self.assertSequenceEqual(free_vars, ["x"])
 
+  @parameterized.named_parameters(
+      ("lambda_1", lambda _x: 3,), ("lambda_2", lambda _x: 3,))
+  def test_multiple_lambda_w_same_line_num_and_args(self, fn):
+    func_map = free_vars_detect._detect_function_free_vars(fn)
+    self.assertEmpty(func_map)
+
   def test_lambda_wo_free_var(self):
     f = lambda x: x + x
     func_map = free_vars_detect._detect_function_free_vars(f)