Gradient for Sum.
(op, grad)
| 136 | |
| 137 | @ops.RegisterGradient("Sum") |
| 138 | def _SumGrad(op, grad): |
| 139 | """Gradient for Sum.""" |
| 140 | # Fast path for when reducing to a scalar and ndims is known: adds only |
| 141 | # Reshape and Tile ops (and possibly a Shape). |
| 142 | input_0_shape = op.inputs[0]._shape_tuple() # pylint: disable=protected-access |
| 143 | if input_0_shape is not None: |
| 144 | axes = tensor_util.constant_value(op.inputs[1]) |
| 145 | if axes is not None: |
| 146 | rank = len(input_0_shape) |
| 147 | if np.array_equal(axes, np.arange(rank)): # Reduce all dims. |
| 148 | if context.executing_eagerly(): |
| 149 | ctx = context.context() |
| 150 | new_shape = ctx.ones_rank_cache().get(rank) |
| 151 | if new_shape is None: |
| 152 | new_shape = constant_op.constant([1] * rank, dtype=dtypes.int32) |
| 153 | ctx.ones_rank_cache().put(rank, new_shape) |
| 154 | else: |
| 155 | new_shape = [1] * rank |
| 156 | grad = array_ops.reshape(grad, new_shape) |
| 157 | # If shape is not fully defined (but rank is), we use Shape. |
| 158 | if None not in input_0_shape: |
| 159 | input_shape = constant_op.constant(input_0_shape, dtype=dtypes.int32) |
| 160 | else: |
| 161 | input_shape = array_ops.shape(op.inputs[0]) |
| 162 | return [array_ops.tile(grad, input_shape), None] |
| 163 | elif None not in input_0_shape and not context.executing_eagerly(): |
| 164 | # The shape and reduction indices are statically known, so we use a |
| 165 | # graph-level cache to avoid recomputing `reduced_shape()` for each |
| 166 | # invocation. |
| 167 | graph = ops.get_default_graph() |
| 168 | |
| 169 | # Canonicalize `axes` to be a tuple of indices. The incoming |
| 170 | # value may be a scalar or a vector, and may include negative indices. |
| 171 | axes = tuple(axes.reshape(-1)) |
| 172 | |
| 173 | try: |
| 174 | output_shape_kept_dims, tile_scaling = graph._reduced_shape_cache[ # pylint: disable=protected-access |
| 175 | (input_0_shape, axes)] |
| 176 | except KeyError: |
| 177 | |
| 178 | # Compute and cache `output_shape_kept_dims` and `tile_scaling`. |
| 179 | def EvaluateAsTuple(t): |
| 180 | value = c_api.TF_TryEvaluateConstant_wrapper( |
| 181 | t.graph._c_graph, t._as_tf_output()) # pylint: disable=protected-access |
| 182 | assert value is not None |
| 183 | return tuple(value) |
| 184 | |
| 185 | output_shape_kept_dims = EvaluateAsTuple( |
| 186 | math_ops.reduced_shape(input_0_shape, axes)) |
| 187 | tile_scaling = EvaluateAsTuple( |
| 188 | _safe_shape_div(input_0_shape, output_shape_kept_dims)) |
| 189 | graph._reduced_shape_cache[(input_0_shape, axes)] = ( # pylint:disable=protected-access |
| 190 | output_shape_kept_dims, tile_scaling) |
| 191 | |
| 192 | grad = array_ops.reshape(grad, output_shape_kept_dims) |
| 193 | return [array_ops.tile(grad, tile_scaling), None] |
| 194 | |
| 195 | input_shape = array_ops.shape(op.inputs[0]) |
no test coverage detected