DynamicPartition(scope, original_indices, partitions, num_partitions);

  // Invert these indices with dynamic_stitch to map the incoming
  // gradients to their source inputs.
  // dynamic_stitch(
  //   { [0, 1, 4], [2, 3] },
  //   { [g1, g2, g3], [g4, g5] })
  // = [g1, g2, g4, g5, g3]
  auto reconstructed =
      DynamicStitch(scope, partitioned_indices.outputs, grad_inputs);

  // should be replaced by 1.
  // We use DynamicStitch to do this.

  // input_rank = 4
  auto input_rank = Size(scope, input_shape);

  // Normalize any negative indices in the reduction_axes to positive
  // values.
  auto axes = Mod(scope, Add(scope, reduction_axes, input_rank), input_rank);

  // This [0..input_rank) range of integers is used in DynamicStitch to
  // first copy input_shape to the result.

#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/random/random.h"

namespace tensorflow {
namespace {

using ops::Const;
using ops::DynamicPartition;
using ops::DynamicStitch;
using ops::Placeholder;

class DataFlowGradTest : public ::testing::Test {
 protected:
  DataFlowGradTest() : scope_(Scope::NewRootScope()) {}

  %0 = "tf.DynamicStitch"(%indices, %arg0) : (tensor<2xi32>, tensor<2x2x2xf32>) -> tensor<2x2x2xf32>
  func.return %0 : tensor<2x2x2xf32>
}

// CHECK-LABEL: func @DynamicStitch_dynamic
func.func @DynamicStitch_dynamic(%arg0: tensor<*xi32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
  // CHECK: tf.DynamicStitch
  %0 = "tf.DynamicStitch"(%arg0, %arg1) : (tensor<*xi32>, tensor<*xf32>) -> tensor<*xf32>

 Const 2
 GreaterEqual 1
 LogicalOr 1
cluster 15 size 614
 Add 22
 AddN 41
 BatchMatMulV2 2
 BiasAddGrad 8
 BroadcastGradientArgs 4
 ConcatOffset 1
 ConcatV2 14
 Const 199
 DynamicStitch 1
 FloorDiv 1
 MatMul 20
 Maximum 1
 Mul 74
 NoOp 13
 Reshape 86
 RsqrtGrad 1
 Select 36
 SigmoidGrad 24
 Slice 32
 StridedSlice 2
 Sub 1
 Sum 11
 TanhGrad 17
 Tile 2

    x = tf.dynamic_stitch(condition_indices, partitioned_data)
    # Here x=[1.1, -1., 6.2, 5.3, -1, 8.4], the -1. values remain
    # unchanged.
```

<div style="width:70%; margin:auto; margin-bottom:10px; margin-top:20px;">
<img style="width:100%" src="https://www.tensorflow.org/images/DynamicStitch.png" alt>
</div>
  }];

  let arguments = (ins

  %0 = "tf.DynamicStitch"(%indices, %arg0) : (tensor<2xi32>, tensor<2x2xf32>) -> tensor<2x2xf32>
  func.return %0 : tensor<2x2xf32>
}

// -----

func.func @testDynamicStitch() -> tensor<2x2xf32> {
  // expected-error @+1 {{requires attribute N with value >= 1}}
  %0 = "tf.DynamicStitch"() : () -> (tensor<2x2xf32>)
  func.return %0 : tensor<2x2xf32>
}

    return success();
  }
};

// Lowers DynamicStitch op with constant indices and with static input and
// output shapes using Reshape, UnPack and Pack op.
//
//   %indices0 = "tf.Const"() {value = dense<4> : tensor<i32>}
//   %indices1 = "tf.Const"() {value = dense<[[3, 2], [1, 0]]> :
//   tensor<2x2xi32>} %0 = "tf.DynamicStitch"(%indices0, %indices1, %arg0,
//   %arg1)

  %indices = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32>

  // CHECK: "tf.Const"() <{value = dense<> : tensor<0x2xf32>}> : () -> tensor<0x2xf32>
  %0 = "tf.DynamicStitch"(%indices, %arg0) : (tensor<0xi32>, tensor<0x2xf32>) -> tensor<0x2xf32>
  func.return %0 : tensor<0x2xf32>
}

// Verifies that tf.Yield op which has no result and is not side effecting is
// preserved.
//

      "DepthwiseConv2dNativeBackpropFilter",
      "DepthwiseConv2dNativeBackpropInput",
      "Dequantize",
      "Diag",
      "DynamicInfeedEnqueueTupleOp",
      "DynamicInfeedDequeueTupleOp",
      "DynamicStitch",
      "DynamicPartition",
      "Einsum",
      "EmptyTensorList",
      "EnsureShape",
      "ExtractImagePatches",
      "Igamma",
      "IgammaGradA",
      "RandomGammaGrad",