Gradient for Prod.
(op, grad)
| 253 | |
| 254 | @ops.RegisterGradient("Prod") |
| 255 | def _ProdGrad(op, grad): |
| 256 | """Gradient for Prod.""" |
| 257 | # The gradient can be expressed by dividing the product by each entry of the |
| 258 | # input tensor, but this approach can't deal with zeros in the input. |
| 259 | # Here, we avoid this problem by composing the output as a product of two |
| 260 | # cumprod operations. |
| 261 | |
| 262 | input_shape = array_ops.shape(op.inputs[0]) |
| 263 | # Reshape reduction indices for the case where the parameter is a scalar |
| 264 | reduction_indices = array_ops.reshape(op.inputs[1], [-1]) |
| 265 | |
| 266 | # Expand grad to full input shape |
| 267 | output_shape_kept_dims = math_ops.reduced_shape(input_shape, op.inputs[1]) |
| 268 | tile_scaling = _safe_shape_div(input_shape, output_shape_kept_dims) |
| 269 | grad = array_ops.reshape(grad, output_shape_kept_dims) |
| 270 | grad = array_ops.tile(grad, tile_scaling) |
| 271 | |
| 272 | # Pack all reduced dimensions into a single one, so we can perform the |
| 273 | # cumprod ops. If the reduction dims list is empty, it defaults to float32, |
| 274 | # so we need to cast here. We put all the shape-related ops on CPU to avoid |
| 275 | # copying back and forth, and since listdiff is CPU only. |
| 276 | with ops.device("/cpu:0"): |
| 277 | rank = array_ops.rank(op.inputs[0]) |
| 278 | reduction_indices = (reduction_indices + rank) % rank |
| 279 | reduced = math_ops.cast(reduction_indices, dtypes.int32) |
| 280 | idx = math_ops.range(0, rank) |
| 281 | other, _ = array_ops.setdiff1d(idx, reduced) |
| 282 | perm = array_ops.concat([reduced, other], 0) |
| 283 | reduced_num = math_ops.reduce_prod(array_ops.gather(input_shape, reduced)) |
| 284 | other_num = math_ops.reduce_prod(array_ops.gather(input_shape, other)) |
| 285 | permuted = array_ops.transpose(op.inputs[0], perm) |
| 286 | permuted_shape = array_ops.shape(permuted) |
| 287 | reshaped = array_ops.reshape(permuted, (reduced_num, other_num)) |
| 288 | |
| 289 | # Calculate product, leaving out the current entry |
| 290 | left = math_ops.cumprod(reshaped, axis=0, exclusive=True) |
| 291 | right = math_ops.cumprod(reshaped, axis=0, exclusive=True, reverse=True) |
| 292 | # For complex inputs, the gradient is in the conjugate direction. |
| 293 | y = array_ops.reshape( |
| 294 | math_ops.conj(left) * math_ops.conj(right), permuted_shape) |
| 295 | |
| 296 | # Invert the transpose and reshape operations. |
| 297 | # Make sure to set the statically known shape information through a reshape. |
| 298 | out = grad * array_ops.transpose(y, array_ops.invert_permutation(perm)) |
| 299 | return array_ops.reshape(out, input_shape), None |
| 300 | |
| 301 | |
| 302 | @ops.RegisterGradient("SegmentSum") |