Gradient for UnsortedSegmentProd. The gradient can be expressed for each segment by dividing the segment's product by each element of the segment input tensor, but this approach can't deal with zeros in the input. Unlike reduce_prod we can't use cumsum here as individual segments may have
(op, grad)
| 463 | |
| 464 | @ops.RegisterGradient("UnsortedSegmentProd") |
| 465 | def _UnsortedSegmentProdGrad(op, grad): |
| 466 | """ Gradient for UnsortedSegmentProd. |
| 467 | |
| 468 | The gradient can be expressed for each segment by dividing the segment's |
| 469 | product by each element of the segment input tensor, but this approach can't |
| 470 | deal with zeros in the input. |
| 471 | Unlike reduce_prod we can't use cumsum here as individual segments may have |
| 472 | a different number of elements. Therefore we consider three cases: |
| 473 | 1) A segment input contains no zeros and we can safely divide by the input |
| 474 | tensor. |
| 475 | 2) A segment contains exactly one zero. Then the gradient of each input of |
| 476 | the segment is zero except for the 0-input, there the gradient is |
| 477 | the product of the remaining segment entries. |
| 478 | 3) A segment contains at least two zeros. The gradient is zero for all |
| 479 | segment inputs. |
| 480 | """ |
| 481 | # Note that unsorted_segment_sum will filter out the negative indices, |
| 482 | # so we don't need to do a logical_and with is_positive here |
| 483 | is_zero = math_ops.equal(op.inputs[0], 0) |
| 484 | num_zeros = gen_math_ops.unsorted_segment_sum( |
| 485 | math_ops.cast(is_zero, dtype=dtypes.int32), op.inputs[1], op.inputs[2]) |
| 486 | # handle case 3 and set the gradient to 0 for segments with more than one |
| 487 | # 0 as input |
| 488 | grad = array_ops.where( |
| 489 | math_ops.greater(num_zeros, 1), array_ops.zeros_like(grad), grad) |
| 490 | # replace all zeros with ones and compute the unsorted_segment_prod |
| 491 | non_zero_data = array_ops.where(is_zero, array_ops.ones_like(op.inputs[0]), |
| 492 | op.inputs[0]) |
| 493 | non_zero_prod = gen_math_ops.unsorted_segment_prod(non_zero_data, |
| 494 | op.inputs[1], op.inputs[2]) |
| 495 | # clip the indices for gather to be positive |
| 496 | zero_clipped_indices = math_ops.maximum(op.inputs[1], |
| 497 | array_ops.zeros_like(op.inputs[1])) |
| 498 | gathered_prod = array_ops.gather(op.outputs[0], zero_clipped_indices) |
| 499 | gathered_non_zero_prod = array_ops.gather(non_zero_prod, zero_clipped_indices) |
| 500 | prod_divided_by_el = gathered_prod / op.inputs[0] # May contain nan/inf. |
| 501 | # Now fetch the individual results for segments containing 0 and those that |
| 502 | # don't. is_zero will also fetch results for entries with negative index |
| 503 | # but the following gather_drop_negatives sets the corresponding entry in |
| 504 | # grad to 0 for these |
| 505 | partial_derivative = array_ops.where(is_zero, gathered_non_zero_prod, |
| 506 | prod_divided_by_el) |
| 507 | gathered_grad = _GatherDropNegatives(grad, op.inputs[1], |
| 508 | zero_clipped_indices)[0] |
| 509 | return gathered_grad * partial_derivative, None, None |
| 510 | |
| 511 | |
| 512 | @ops.RegisterGradient("Abs") |