data_flow_grad.py revision a558c6e3b38846727873b5afbbc3ba309ae5dff5
1# Copyright 2015 The TensorFlow Authors. 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 22 23from tensorflow.python.framework import dtypes 24from tensorflow.python.framework import ops 25from tensorflow.python.ops import array_ops 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("QueueDequeueUpTo") 72ops.NoGradient("QueueClose") 73ops.NoGradient("QueueSize") 74 75ops.NoGradient("Stack") 76ops.NoGradient("StackPush") 77ops.NoGradient("StackPop") 78ops.NoGradient("StackClose") 79 80ops.NoGradient("GetSessionHandle") 81ops.NoGradient("GetSessionTensor") 82ops.NoGradient("DeleteSessionTensor") 83