Optimized version of `broadcast_gradient_args` that caches results. This implementation avoids creating `broadcast_gradient_args` ops in the case that the input shapes are fully defined, and provides hints to the calling code that can be used to avoid creating reduction and reshaping ops.
(x, y, grad)
| 55 | |
| 56 | |
| 57 | def SmartBroadcastGradientArgs(x, y, grad): |
| 58 | """Optimized version of `broadcast_gradient_args` that caches results. |
| 59 | |
| 60 | This implementation avoids creating `broadcast_gradient_args` ops in the case |
| 61 | that the input shapes are fully defined, and provides hints to the calling |
| 62 | code that can be used to avoid creating reduction and reshaping ops. |
| 63 | |
| 64 | Args: |
| 65 | x: The left input tensor to a broadcasting binary op. |
| 66 | y: The right input tensor to a broadcasting binary op. |
| 67 | grad: The incoming gradient tensor for a broadcasting binary op. |
| 68 | |
| 69 | Returns: |
| 70 | A pair of tuples, containing: |
| 71 | * A 3-tuple of broadcast information for x, containing: |
| 72 | * The shape of x (as a tuple or Tensor). |
| 73 | * The reduction indices for x (as a tuple or Tensor). |
| 74 | * A boolean, which if True, indicates that x's shape differs from grad's |
| 75 | shape (and so x's gradient must be reduced and/or reshaped). |
| 76 | * A 3-tuple of broadcast information for y, containing the respective |
| 77 | details for y. |
| 78 | """ |
| 79 | # NOTE: It may be productive to apply these optimizations in the eager case |
| 80 | # as well. |
| 81 | if context.executing_eagerly() or not ( |
| 82 | isinstance(x, ops.Tensor) and isinstance(y, ops.Tensor) |
| 83 | and isinstance(grad, ops.Tensor)): |
| 84 | sx = array_ops.shape(x) |
| 85 | sy = array_ops.shape(y) |
| 86 | rx, ry = gen_array_ops.broadcast_gradient_args(sx, sy) |
| 87 | return (sx, rx, True), (sy, ry, True) |
| 88 | |
| 89 | # pylint: disable=protected-access |
| 90 | x_shape_tuple = x._shape_tuple() |
| 91 | y_shape_tuple = y._shape_tuple() |
| 92 | grad_shape_tuple = grad._shape_tuple() |
| 93 | # pylint: enable=protected-access |
| 94 | |
| 95 | if (x_shape_tuple is None or None in x_shape_tuple or |
| 96 | y_shape_tuple is None or None in y_shape_tuple): |
| 97 | sx = array_ops.shape_internal(x, optimize=False) |
| 98 | sy = array_ops.shape_internal(y, optimize=False) |
| 99 | rx, ry = gen_array_ops.broadcast_gradient_args(sx, sy) |
| 100 | return (sx, rx, True), (sy, ry, True) |
| 101 | |
| 102 | x_needs_reduction = x_shape_tuple != grad_shape_tuple |
| 103 | y_needs_reduction = y_shape_tuple != grad_shape_tuple |
| 104 | |
| 105 | # Get the default graph rather than relying on `x.graph`, `y.graph`, or |
| 106 | # `grad.graph`, because these may be eager tensors. |
| 107 | g = ops.get_default_graph() |
| 108 | |
| 109 | try: |
| 110 | rx, ry = g._bcast_grad_args_cache[(x_shape_tuple, y_shape_tuple)] # pylint: disable=protected-access |
| 111 | return (x_shape_tuple, rx, x_needs_reduction), ( |
| 112 | y_shape_tuple, ry, y_needs_reduction) |
| 113 | except KeyError: |
| 114 | rx, ry = array_ops.broadcast_gradient_args(x_shape_tuple, y_shape_tuple) |
no test coverage detected