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

1# Copyright 2019 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"""Utilities for managing forward accumulators. 

16 

17A separate file from forwardprop.py so that functions can use these utilities. 

18""" 

19 

20import collections 

21import contextlib 

22 

23from tensorflow.python import pywrap_tfe 

24 

25 

26class TangentInfo( 

27 collections.namedtuple("TangentInfo", ["indices", "tangents"])): 

28 """Packed forward accumulator state. The return value of `pack_tangents`.""" 

29 

30 def __new__(cls, indices=None, tangents=None): 

31 if indices is None: 

32 indices = () 

33 if tangents is None: 

34 tangents = [] 

35 return super(TangentInfo, cls).__new__(cls, indices, tangents) 

36 

37 

38def pack_tangents(tensors): 

39 """Packs forward accumulator state into a TangentInfo tuple. 

40 

41 Args: 

42 tensors: A flat list of Tensors to pack forward accumulator state for. 

43 

44 Returns: 

45 A tuple of (indices, tangents): 

46 indices: A sequence of sequences of two-element tuples. Each forward 

47 accumulator is represented as a sequence of tuples with (primal_index, 

48 jvp_index). Both integers index into the concatenated `tensors + jvps` 

49 array. 

50 tangents: A flat list of Tensors. Best interpreted as a sequence to be 

51 appended to `tensors`. 

52 """ 

53 return TangentInfo(*pywrap_tfe.TFE_Py_PackJVPs(tensors)) 

54 

55 

56@contextlib.contextmanager 

57def push_forwardprop_state(): 

58 """Temporarily push or pop transient state for accumulators in the active set. 

59 

60 Allows an accumulator which is currently processing an operation to 

61 temporarily reset its state. This is useful when building forwardprop versions 

62 of functions, where an accumulator will trigger function building and then 

63 must process captured symbolic tensors while building it. Without pushing and 

64 popping, accumulators ignore operations executed as a direct result of their 

65 own jvp computations. 

66 

67 Yields: 

68 None (used for its side effect). 

69 """ 

70 try: 

71 pywrap_tfe.TFE_Py_ForwardAccumulatorPushState() 

72 yield 

73 finally: 

74 pywrap_tfe.TFE_Py_ForwardAccumulatorPopState()