auto result_type =
      OpTrait::util::getBroadcastedType(x.getType(), y.getType());
  if (!result_type) {
    if (incompatible_shape_error.getValue()) {
      mlir::emitError(loc, "non-broadcastable operands");
    } else {
      return UnrankedTensorType::get(builder->getI1Type());
    }
  }

  auto ranked_type = mlir::dyn_cast<RankedTensorType>(result_type);

    // Note that we are using the same `result_type` for both the
    // `TensorListStackOp` and `ConcatOp`, since the first dimension of the
    // shape specified by `result_type` is -1.
    auto stacked_extended_part = rewriter->create<TF::TensorListStackOp>(
        loc, result_type, extended_part,
        /*element_shape=*/CreateI32SplatConst(loc, rewriter, {}, -1),

  // All the rest: states, device.
  for (int i = 2; i < 5; ++i) {
    auto result_type =
        mlir::dyn_cast<RankedTensorType>(func_op.getResultTypes()[i]);
    outputs.push_back(CreatTfF32ConstOp(builder, result_type.getShape(), 0.0f,
                                        func_op.getLoc()));
    output_types.push_back(result_type);
  }

  // Update function signatures.

  if (x_type.getShape() != y_type.getShape()) return failure();

  auto result_type = squared_diff_op.getType();
  if (!result_type) return failure();

  auto sub_op =
      rewriter.create<TF::SubOp>(squared_diff_op.getLoc(), result_type, x, y);
  auto mul_op =
      rewriter.create<TF::MulOp>(squared_diff_op.getLoc(), result_type,
                                 sub_op.getResult(), sub_op.getResult());

      TensorType result_type = mlir::dyn_cast<TensorType>(result.getType());
      if (!result_type) {
        result_types.push_back(result.getType());
        continue;
      }
      QuantizedType qtype =
          mlir::dyn_cast<QuantizedType>(result_type.getElementType());
      if (!qtype) {
        result_types.push_back(result_type);
        continue;
      }

    const auto result_type = mlir::cast<TensorType>(op.getResult(0).getType());
    const SmallVector<int64_t> new_result_shape =
        Permute<int64_t>(result_type.getShape(), kNchwToNhwcPermutation);

    const TensorType new_result_type =
        result_type.cloneWith(new_result_shape, result_type.getElementType());

  auto matmul = rewriter.create<TFL::BatchMatMulOp>(
      loc, RankedTensorType::get(matmul_shape, result_type.getElementType()),
      lhs_flattend, rhs_flattend, /*adj_x*/ false_attr, /*adj_y*/ false_attr,
      /*asym_quant_input*/ false_attr);
  if (result_type.hasStaticShape()) {
    auto reshaped =
        rewriter.create<mhlo::ReshapeOp>(loc, result_type, matmul.getResult());
    return reshaped.getResult();
  }

  values.reserve(rank);
  for (int i = 0; i < rank; ++i) values.push_back(APInt(bitwidth, r1[i]));
  auto result_type = tensorflow::GetTypeFromTFTensorShape(
      {rank}, IntegerType::get(builder.getContext(), bitwidth));
  return builder.create<TF::ConstOp>(
      loc, DenseElementsAttr::get(result_type, values));
}

Value GetIndicesForElement(Value index, Value buffer, OpBuilder builder,

    // Verify that the first result type is same as the rest of the results.
    // We skip the comparison against itself.
    for (auto result_type : llvm::drop_begin(op->getResultTypes(), 1)) {
      if (!mlir::tf_type::HasCompatibleElementTypes(type, result_type))
        return op->emitOpError()
               << "requires all return types to have compatible element types";
    }

    auto result_type =
        mlir::dyn_cast_or_null<ShapedType>(transpose_op.getResult().getType());
    auto is_valid_move =
        llvm::all_of(op->getOperands(), [result_type](Value operand) -> bool {
          auto operand_type =
              mlir::dyn_cast_or_null<ShapedType>(operand.getType());
          return result_type && operand_type && result_type.hasRank() &&