message: speedup reduce op grads when keep_dims=True
diff: @@ -161,22 +161,26 @@ def _SumGrad(op, grad):
  return [array_ops.tile(grad, input_shape), None]
  input_shape = array_ops.shape(op.inputs[0])
  +
  + if not op.get_attr("keep_dims"):
  # TODO(apassos) remove this once device placement for eager ops makes more
  # sense.
  with ops.colocate_with(input_shape):
  output_shape_kept_dims = math_ops.reduced_shape(input_shape, op.inputs[1])
  - tile_scaling = _safe_shape_div(input_shape, output_shape_kept_dims)
  grad = array_ops.reshape(grad, output_shape_kept_dims)
  - return [array_ops.tile(grad, tile_scaling), None]
  + return [array_ops.broadcast_to(grad, input_shape), None]
  def _MinOrMaxGrad(op, grad):
  """Gradient for Min or Max. Amazingly it's precisely the same code."""
  input_shape = array_ops.shape(op.inputs[0])
  - output_shape_kept_dims = math_ops.reduced_shape(input_shape, op.inputs[1])
  y = op.outputs[0]
  + if not op.get_attr("keep_dims"):
  + output_shape_kept_dims = math_ops.reduced_shape(input_shape, op.inputs[1])
  y = array_ops.reshape(y, output_shape_kept_dims)
  grad = array_ops.reshape(grad, output_shape_kept_dims)
  + else:
  + output_shape_kept_dims = array_ops.shape(y)
  # Compute the number of selected (maximum or minimum) elements in each
  # reduction dimension. If there are multiple minimum or maximum elements
  @@ -232,10 +236,11 @@ def _ProdGrad(op, grad):
  reduction_indices = array_ops.reshape(op.inputs[1], [-1])
  # Expand grad to full input shape
  + if not op.get_attr("keep_dims"):
  output_shape_kept_dims = math_ops.reduced_shape(input_shape, op.inputs[1])
  - tile_scaling = _safe_shape_div(input_shape, output_shape_kept_dims)
  grad = array_ops.reshape(grad, output_shape_kept_dims)
  - grad = array_ops.tile(grad, tile_scaling)
  +
  + grad = array_ops.broadcast_to(grad, input_shape)
  # Pack all reduced dimensions into a single one, so we can perform the
  # cumprod ops. If the reduction dims list is empty, it defaults to float32,
  