Get the aggregated gradients for op. Args: grads: The map of memoized gradients. op: The op to get gradients for. gradient_uid: A unique identifier within the graph indicating which invocation of gradients is being executed. Used to cluster ops for compilation. loop_st
(grads,
op,
gradient_uid,
loop_state,
aggregation_method=None)
| 948 | |
| 949 | |
| 950 | def _AggregatedGrads(grads, |
| 951 | op, |
| 952 | gradient_uid, |
| 953 | loop_state, |
| 954 | aggregation_method=None): |
| 955 | """Get the aggregated gradients for op. |
| 956 | |
| 957 | Args: |
| 958 | grads: The map of memoized gradients. |
| 959 | op: The op to get gradients for. |
| 960 | gradient_uid: A unique identifier within the graph indicating |
| 961 | which invocation of gradients is being executed. Used to cluster |
| 962 | ops for compilation. |
| 963 | loop_state: An object for maintaining the state of the while loops in the |
| 964 | graph. It is of type ControlFlowState. None if the graph |
| 965 | contains no while loops. |
| 966 | aggregation_method: Specifies the method used to combine gradient terms. |
| 967 | Accepted values are constants defined in the class `AggregationMethod`. |
| 968 | |
| 969 | Returns: |
| 970 | A list of gradients, one per each output of `op`. If the gradients |
| 971 | for a particular output is a list, this function aggregates it |
| 972 | before returning. |
| 973 | |
| 974 | Raises: |
| 975 | TypeError: if the incoming grads are not Tensors or IndexedSlices. |
| 976 | ValueError: if the arguments are invalid. |
| 977 | |
| 978 | """ |
| 979 | if aggregation_method is None: |
| 980 | aggregation_method = AggregationMethod.DEFAULT |
| 981 | if aggregation_method not in [ |
| 982 | AggregationMethod.ADD_N, AggregationMethod.EXPERIMENTAL_TREE, |
| 983 | AggregationMethod.EXPERIMENTAL_ACCUMULATE_N |
| 984 | ]: |
| 985 | raise ValueError( |
| 986 | "Invalid aggregation_method specified %s." % aggregation_method) |
| 987 | out_grads = _GetGrads(grads, op) |
| 988 | for i, out_grad in enumerate(out_grads): |
| 989 | if loop_state: |
| 990 | if isinstance(out_grad, (ops.Tensor, ops.IndexedSlices)): |
| 991 | assert control_flow_util.IsLoopSwitch(op) |
| 992 | continue |
| 993 | # Grads have to be Tensors or IndexedSlices |
| 994 | if (isinstance(out_grad, collections_abc.Sequence) and not all( |
| 995 | isinstance(g, (ops.Tensor, ops.IndexedSlices)) |
| 996 | for g in out_grad |
| 997 | if g is not None)): |
| 998 | raise TypeError("gradients have to be either all Tensors " |
| 999 | "or all IndexedSlices") |
| 1000 | # Aggregate multiple gradients, and convert [] to None. |
| 1001 | if out_grad: |
| 1002 | if len(out_grad) < 2: |
| 1003 | used = "nop" |
| 1004 | out_grads[i] = out_grad[0] |
| 1005 | elif all(isinstance(g, ops.Tensor) for g in out_grad if g is not None): |
| 1006 | tensor_shape = _AccumulatorShape(out_grad) |
| 1007 | if (aggregation_method == AggregationMethod.EXPERIMENTAL_ACCUMULATE_N |
no test coverage detected