return failure();
      }
    } else if (auto addn = llvm::dyn_cast<TF::AddNOp>(&op)) {
      auto it = buffer_to_size->find(addn.getOperand(0));
      if (it != buffer_to_size->end()) {
        addn.getSum().setType(
            mlir::cast<TensorType>(addn.getOperand(0).getType()));
        auto size = it->getSecond();
        (*buffer_to_size)[addn.getSum()] = size;
      }

      });
  return has_variant_input || has_variant_output;
}

// There are 2 standard tf ops which are not TensorList ops that may take as
// input a tensorlist. These are tf.AddN and tf.ZeroesLike. Since the runtime
// implementation of a tensorlist are not compatible between tf and tflite
// we cannot use tflite tensorlist kernels until these cases are handled.

  // CHECK: [[ADD_N:%.*]] = "tf.AddN"(%arg0, [[ZERO]], [[ONE]])
  // CHECK: return %arg0, %arg0, [[ZERO]], [[ADD_N]]
  %2 = "tf.AddN"(%arg0, %1, %1) : (tensor<2xf32>, tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
  %3 = "tf.AddN"(%1, %arg0, %1) : (tensor<2xf32>, tensor<2xf32> , tensor<2xf32>) -> tensor<2xf32>
  %4 = "tf.AddN"(%1, %1) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>

  auto c = ops::Placeholder(root.WithOpName("C"), DT_FLOAT);
  auto d = ops::Add(root.WithOpName("D"), a, b);
  auto e = ops::Add(root.WithOpName("E"), d, c);
  auto f = ops::Neg(root.WithOpName("F"), e);
  auto g = ops::AddN(root.WithOpName("G"), std::initializer_list<Output>{e, f});

  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
  TF_CHECK_OK(root.ToGraph(graph.get()));

  GraphShapeInfo shape_info;

  return rewriter.create<ConcatV2Op>(loc, type, ValueRange(vals), axis);
}

// Lowers AddN op to a sequence of AddV2 ops to accumulate operands.
//
// Note that to improve the parallelism, AddN op uses tree-based reduction.
// For example, tf.AddN([0, 1, 2, 3, 4]) behaves as follows:
//
//                 0     1     2     3     4
//                 |     |     |     |     |

  // CHECK: [[d:%.*]] = tf_mlrt.await [[futures]]#2
  %d = "tf_mlrt.tf_await"(%future_d) : (!mlrt.future) ->tensor<i32>

  // CHECK: [[result:%.*]] = tf_mlrt.executeop([[b]], [[d]], [[f]])
  // CHECK-SAME: AddN
  %result = "tf.AddN"(%b, %d, %f) {__op_key = 9: i32}: (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<i32>

  // CHECK: mlrt.await_handle [[handle_0]]
  // CHECK: mlrt.await_handle [[handle_1]]
  mlrt.await_handle %handle_0

}

// -----

// CHECK-LABEL: variantAddN
func.func @variantAddN(%arg0: tensor<!tf_type.variant<tensor<*xi32>>>, %arg1: tensor<!tf_type.variant<tensor<*xi32>>>) -> tensor<!tf_type.variant<tensor<*xi32>>> {
  %1 = "tf.AddN"(%arg0, %arg1) : (tensor<!tf_type.variant<tensor<*xi32>>>, tensor<!tf_type.variant<tensor<*xi32>>>) -> tensor<!tf_type.variant<tensor<*xi32>>>

  if (gradient_tensors.size() == 1) {
    return gradient_tensors[0];
  }

  AbstractOperationPtr op(ctx_->CreateOperation());
  Status s = op->Reset("AddN", /*raw_device_name=*/nullptr);
  if (!s.ok()) {
    return nullptr;
  }
  s = op->AddInputList(gradient_tensors);
  if (!s.ok()) {
    return nullptr;
  }

  int num_outputs = 1;

  // CHECK: tf.AddN
  %0 = "tf.AddN"(%arg0, %arg1, %arg2) : (tensor<!tf_type.variant<tensor<2xf32>>>, tensor<!tf_type.variant<tensor<2xf32>>>, tensor<!tf_type.variant<tensor<2xf32>>>) -> tensor<!tf_type.variant<tensor<2xf32>>>
  func.return %0 : tensor<!tf_type.variant<tensor<2xf32>>>

  }

  // CHECK-LABEL: unsupported_dtype
  func.func @unsupported_dtype(%arg0: tensor<2x!tf_type.variant>) -> tensor<2x!tf_type.variant> {
    // CHECK: tf.AddN
    // expected-remark@+1 {{skipping legalization due to unsupported type 'tensor<2x!tf_type.variant>'}}
    %0 = "tf.AddN"(%arg0, %arg0) : (tensor<2x!tf_type.variant>, tensor<2x!tf_type.variant>) -> tensor<2x!tf_type.variant>

    func.return %0 : tensor<2x!tf_type.variant>
  }