Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow/python/eager/polymorphic_function/autograph_util.py: 33%

1# Copyright 2023 The TensorFlow Authors. All Rights Reserved. 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================== 

15# pylint: disable=unidiomatic-typecheck 

16"""Autograph utility functions for polymorphic_function.""" 

17 

18from tensorflow.python.autograph.core import converter 

19from tensorflow.python.autograph.impl import api 

20from tensorflow.python.util import tf_decorator 

21 

22 

23def py_func_from_autograph( 

24 python_func, 

25 autograph_options=None, 

26): 

27 """Compile a python function using autograph, for use with FuncGraph. 

28 

29 Args: 

30 python_func: the Python function to compile. 

31 autograph_options: additional knobs to control when `autograph=True`. 

32 See https://www.tensorflow.org/guide/autograph for more information. 

33 Returns: 

34 python_func, converted using autograph. 

35 """ 

36 _, original_func = tf_decorator.unwrap(python_func) 

37 

38 def autograph_handler(*args, **kwargs): 

39 """Calls a converted version of original_func.""" 

40 try: 

41 return api.converted_call( 

42 original_func, 

43 args, 

44 kwargs, 

45 options=converter.ConversionOptions( 

46 recursive=True, 

47 optional_features=autograph_options, 

48 user_requested=True, 

49 )) 

50 except Exception as e: # pylint:disable=broad-except 

51 if hasattr(e, "ag_error_metadata"): 

52 raise e.ag_error_metadata.to_exception(e) 

53 else: 

54 raise 

55 

56 # Wrapping around a decorator allows checks like tf_inspect.getargspec 

57 # to be accurate. 

58 converted_func = tf_decorator.make_decorator(original_func, autograph_handler) 

59 return tf_decorator.rewrap(python_func, original_func, converted_func)