| 42 | REGISTER_NO_GRADIENT_OP("DeleteSessionTensor"); |
| 43 | |
| 44 | Status DynamicPartitionGrad(const Scope& scope, const Operation& op, |
| 45 | const std::vector<Output>& grad_inputs, |
| 46 | std::vector<Output>* grad_outputs) { |
| 47 | // DynamicPartition only moves input values into various positions |
| 48 | // in the output, so the gradient operation only has to map incoming |
| 49 | // gradients into their input source locations. |
| 50 | // running example: |
| 51 | // data = [10, 20, 30, 40, 50] |
| 52 | // partitions = [0, 0, 1, 1, 0] |
| 53 | // num_partitions = 2 |
| 54 | // dynamic_partition(data, partitions, num_partitions) = { |
| 55 | // [10, 20, 50], |
| 56 | // [30, 40] |
| 57 | // } |
| 58 | // grads = { |
| 59 | // [g1, g2, g3], |
| 60 | // [g4, g5] |
| 61 | // } |
| 62 | // The desired propagation of the gradients back to the data inputs is: |
| 63 | // [g1, g2, g4, g5, g3] |
| 64 | auto data = op.input(0); |
| 65 | auto partitions = op.input(1); |
| 66 | int32 num_partitions; |
| 67 | TF_RETURN_IF_ERROR( |
| 68 | GetNodeAttr(op.node()->attrs(), "num_partitions", &num_partitions)); |
| 69 | |
| 70 | // Note: the shape of the partitions is a prefix of the data shape. |
| 71 | // shape(partitions) = [5] |
| 72 | auto partitions_shape = Shape(scope, partitions); |
| 73 | // We now create a partitions-shaped tensor with integers from |
| 74 | // [0..size(partitions)) This will be dynamic_partitioned with the |
| 75 | // input parameters, providing the destination index for a given |
| 76 | // source item. |
| 77 | // partitions_size = prod([5]) = 5 |
| 78 | // reshape(range(partitions_size), [5]) = [0, 1, 2, 3, 4] |
| 79 | auto zero = Const(scope, 0); |
| 80 | auto one = Const(scope, 1); |
| 81 | auto original_indices = Reshape( |
| 82 | scope, Range(scope, zero, Prod(scope, partitions_shape, zero), one), |
| 83 | partitions_shape); |
| 84 | // dynamic_partition( |
| 85 | // [0, 1, 2, 3, 4], |
| 86 | // [0, 0, 1, 1, 0], 2) |
| 87 | // = { [0, 1, 4], |
| 88 | // [2, 3] } |
| 89 | auto partitioned_indices = |
| 90 | DynamicPartition(scope, original_indices, partitions, num_partitions); |
| 91 | |
| 92 | // Invert these indices with dynamic_stitch to map the incoming |
| 93 | // gradients to their source inputs. |
| 94 | // dynamic_stitch( |
| 95 | // { [0, 1, 4], [2, 3] }, |
| 96 | // { [g1, g2, g3], [g4, g5] }) |
| 97 | // = [g1, g2, g4, g5, g3] |
| 98 | auto reconstructed = |
| 99 | DynamicStitch(scope, partitioned_indices.outputs, grad_inputs); |
| 100 | // reshape back into a data-shaped tensor to propagate gradients for the data |
| 101 | // input. |
nothing calls this directly
no test coverage detected