data_flow_grad.py revision e14090d217fc4e7e49ac04ccbc50acdba8b9f120
1# Copyright 2015 Google Inc. All Rights Reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15 16"""Gradients for operators defined in data_flow_ops.py.""" 17from __future__ import absolute_import 18from __future__ import division 19from __future__ import print_function 20 21from six.moves import xrange # pylint: disable=redefined-builtin 22from tensorflow.python.framework import dtypes 23from tensorflow.python.framework import ops 24from tensorflow.python.ops import array_ops 25from tensorflow.python.ops import constant_op 26from tensorflow.python.ops import data_flow_ops 27from tensorflow.python.ops import math_ops 28 29 30@ops.RegisterGradient("DynamicPartition") 31def _DynamicPartitionGrads(op, *grads): 32 """Gradients for DynamicPartition.""" 33 data = op.inputs[0] 34 indices = op.inputs[1] 35 num_partitions = op.get_attr("num_partitions") 36 37 prefix_shape = array_ops.shape(indices) 38 original_indices = array_ops.reshape( 39 math_ops.range(math_ops.reduce_prod(prefix_shape)), prefix_shape) 40 partitioned_indices = data_flow_ops.dynamic_partition( 41 original_indices, indices, num_partitions) 42 reconstructed = data_flow_ops.dynamic_stitch(partitioned_indices, grads) 43 reconstructed = array_ops.reshape(reconstructed, array_ops.shape(data)) 44 return [reconstructed, None] 45 46 47@ops.RegisterGradient("DynamicStitch") 48def _DynamicStitchGrads(op, grad): 49 """Gradients for DynamicStitch.""" 50 51 num_values = len(op.inputs) // 2 52 indices_grad = [None] * num_values 53 54 def AsInt32(x): 55 return (x if op.inputs[0].dtype == dtypes.int32 else 56 math_ops.cast(x, dtypes.int32)) 57 inputs = [AsInt32(op.inputs[i]) for i in xrange(num_values)] 58 if isinstance(grad, ops.IndexedSlices): 59 output_shape = array_ops.shape(op.outputs[0]) 60 output_rows = output_shape[0] 61 grad = math_ops.unsorted_segment_sum(grad.values, grad.indices, output_rows) 62 values_grad = [array_ops.gather(grad, inp) for inp in inputs] 63 return indices_grad + values_grad 64 65 66ops.NoGradient("Queue") 67ops.NoGradient("QueueEnqueue") 68ops.NoGradient("QueueEnqueueMany") 69ops.NoGradient("QueueDequeue") 70ops.NoGradient("QueueDequeueMany") 71ops.NoGradient("QueueClose") 72ops.NoGradient("QueueSize") 73 74ops.NoGradient("Stack") 75ops.NoGradient("StackPush") 76ops.NoGradient("StackPop") 77ops.NoGradient("StackClose") 78 79ops.NoGradient("TensorArray") 80ops.NoGradient("TensorArrayGrad") 81ops.NoGradient("TensorArrayClose") 82 83 84@ops.RegisterGradient("TensorArrayRead") 85def _TensorArrayReadGrad(op, grad): 86 handle = op.inputs[0] 87 index = op.inputs[1] 88 dtype = op.get_attr("dtype") 89 g = data_flow_ops.TensorArray(size=None, dtype=dtype, handle=handle).grad() 90 w_g = g.write(index, grad) 91 return [None, None, w_g.flow] 92 93 94@ops.RegisterGradient("TensorArrayWrite") 95def _TensorArrayWriteGrad(op, flow): 96 # handle is the output store_handle of TensorArrayReadGrad or 97 # the handle output of TensorArrayWriteGrad. we must use this one. 98 handle = op.inputs[0] 99 index = op.inputs[1] 100 dtype = op.get_attr("T") 101 g = data_flow_ops.TensorArray(size=None, dtype=dtype, handle=handle).grad() 102 with ops.control_dependencies([flow]): 103 grad = g.read(index) 104 return [None, None, grad, flow] 105 106 107@ops.RegisterGradient("TensorArrayPack") 108def _TensorArrayPackGrad(op, grad): 109 handle = op.inputs[0] 110 dtype = op.get_attr("dtype") 111 g = data_flow_ops.TensorArray(size=None, dtype=dtype, handle=handle).grad() 112 u_g = g.unpack(grad) 113 return [None, u_g.flow] 114 115 116@ops.RegisterGradient("TensorArrayUnpack") 117def _TensorArrayUnpackGrad(op, flow): 118 # handle is the output store_handle of TensorArrayReadGrad or 119 # the handle output of TensorArrayUnpackGrad. we must use this one. 120 handle = op.inputs[0] 121 dtype = op.get_attr("T") 122 g = data_flow_ops.TensorArray(size=None, dtype=dtype, handle=handle).grad() 123 with ops.control_dependencies([flow]): 124 grad = g.pack() 125 return [None, grad, flow] 126# pylint: enable=protected-access 127