result_types.push_back(op->getOpResult(2).getType());

    // Populate attributes.
    SmallVector<NamedAttribute, 4> attributes;
    // Activation will always be tanh.
    attributes.push_back(rewriter.getNamedAttr("fused_activation_function",
                                               rewriter.getStringAttr("TANH")));
    // cell_clip.
    attributes.push_back(

  const std::string gpu1 = "/job:worker/replica:1/task:2/device:GPU:1";

  llvm::SmallVector<std::unique_ptr<Device>, 2> devices;
  devices.push_back(FakeDevice::Make(cpu0));
  devices.push_back(FakeDevice::Make(gpu0, "compute capability: 7.0"));
  devices.push_back(FakeDevice::Make(gpu1));

  DeviceSet device_set;
  for (auto& device : devices) device_set.AddDevice(device.get());
  AddDevicesToOp(*module_ref, &device_set);

    mlir::TensorType mlir_ty = mlir::UnrankedTensorType::get(elt_ty);
    input_tys.push_back(mlir_ty);
  }
  for (auto ty : output_dtys) {
    mlir::Type elt_ty;
    TF_RETURN_IF_ERROR(ConvertDataType(ty, builder, &elt_ty));
    mlir::TensorType mlir_ty = mlir::UnrankedTensorType::get(elt_ty);
    output_tys.push_back(mlir_ty);
  }
  llvm::SmallVector<mlir::NamedAttribute, 4> attrs;

  std::vector<std::string> disabled_checks;
  for (auto attr : op.getDisabledChecks().getAsRange<StringAttr>()) {
    disabled_checks.push_back(attr.getValue().str());
  }
  std::vector<std::string> platforms;
  for (auto attr : op.getPlatforms().getAsRange<StringAttr>()) {
    platforms.push_back(attr.getValue().str());
  }
  TF_ASSIGN_OR_RETURN(
      auto loader,
      tensorflow::XlaCallModuleLoader::Create(

  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());

  std::vector<TensorHandlePtr> vector_handles;
  vector_handles.reserve(2);
  vector_handles.push_back(std::move(two_vector));
  vector_handles.push_back(std::move(three_vector));
  std::unique_ptr<ParallelTensor> unknown_length_vector =
      ParallelTensor::FromTensorHandles(
          parallel_device, std::move(vector_handles), status.get());

  } else {
    SmallVector<double> scales_vec;
    SmallVector<int64_t> zero_points_vec;
    for (auto elem : scales.getValues<float>()) scales_vec.push_back(elem);
    for (auto elem : zero_points.getValues<int32_t>())
      zero_points_vec.push_back(elem);
    elem_ty = quant::UniformQuantizedPerAxisType::get(
        flags, storage_type, expressed_type, scales_vec, zero_points_vec,

                      arg_locs);
  SmallVector<NamedAttribute> func_attrs;
  func_attrs.push_back(
      {StringAttr::get(context, "inputs"),
       StringAttr::get(context, absl::StrJoin(input_names, ","))});
  func_attrs.push_back(
      {StringAttr::get(context, "outputs"),
       StringAttr::get(context, absl::StrJoin(output_names, ","))});

    return split.emitOpError("different split lengths are not supported");
  }
  llvm::SmallVector<int64_t, 8> elem_shape;
  elem_shape.push_back(length);
  for (int64_t dim : buffer_type.getShape().drop_front()) {
    elem_shape.push_back(dim);
  }
  *elem_type = RankedTensorType::get(elem_shape, buffer_type.getElementType());
  return success();
}

// Tries to infer the tensor array element shape.

}

TEST_F(ArrayGradTest, ConcatV2Grad) {
  TensorShape shape({3, 2, 5});
  std::vector<Output> xs;
  xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape)));
  xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape)));
  xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape)));
  auto axis = Const(scope_, 0);
  auto y = Concat(scope_, xs, axis);

    func::FuncOp func,
    llvm::function_ref<WalkResult(func::FuncOp)> callback,
    SymbolTableCollection* symbol_table) {
  llvm::SmallDenseSet<Operation*> visited;

  llvm::SmallVector<func::FuncOp> stack;
  stack.push_back(func);

  while (!stack.empty()) {
    func::FuncOp f = stack.back();
    stack.pop_back();

    if (!visited.insert(f).second) {
      continue;
    }

    WalkResult result = callback(f);