}
}

// Returns a Constant op with a splat vector value.
absl::StatusOr<arith::ConstantOp> CreateConstOpWithVectorValue(
    PatternRewriter* rewriter, Location loc, ShapedType shaped_type,
    int value) {
  ShapedType dense_type = RankedTensorType::get(shaped_type.getShape(),
                                                shaped_type.getElementType());
  auto attr = CreateTypedAttr(dense_type, value);

  dot_dimension_numbers.ParseFromString(dot_dimension_numbers_str.str());
  SmallVector<Value> input_arguments = {lhs, rhs};
  const int lhs_rank = mlir::cast<ShapedType>(lhs.getType()).getShape().size();
  const int rhs_rank = mlir::cast<ShapedType>(rhs.getType()).getShape().size();

  const std::string einsum_equation =
      CreateEinsumEquation(dot_dimension_numbers, lhs_rank, rhs_rank);

  auto lhs = dyn_cast_or_null<DenseElementsAttr>(lhs_op.getValue());
  auto rhs = dyn_cast_or_null<DenseElementsAttr>(rhs_op.getValue());
  if (!lhs || !rhs) return {};

  ShapedType type = mlir::cast<ShapedType>(op->getType());
  if (!type.hasStaticShape()) {
    return {};
  }

  Type etype = type.getElementType();

  // Evaluate for element types.
  if (!mlir::isa<ElementType>(etype)) {

  CPred<"IsLastDimensionEqualOne($0) || "
        "(!$0.getType().cast<ShapedType>().hasStaticShape() && "
        "  $0.getType().cast<ShapedType>().hasRank() && "
        "  $0.getType().cast<ShapedType>().getRank() == 2 && "
        "  !$0.getType().cast<ShapedType>().getShape().empty() && "
        "  $0.getType().cast<ShapedType>().getShape()[1] == 1)">>;

// Replace
//   Equal(X, indices)
// With

absl::StatusOr<arith::ConstantOp> CreateConstOpWithSingleValue(
    PatternRewriter* rewriter, Location loc, ShapedType shaped_type, int value);

// Returns a Constant op with a splat vector value.
absl::StatusOr<arith::ConstantOp> CreateConstOpWithVectorValue(
    PatternRewriter* rewriter, Location loc, ShapedType shaped_type, int value);

}  // namespace TFL
}  // namespace mlir

//   bool hasRank()
//   ArrayRef<int64_t> getShape()
// This includes mlir::TF::ShapeAttr and mlir::ShapedType.
template <typename ShapeContainerT>
void SetTensorShapeProto(ShapeContainerT shape, TensorShapeProto* proto) {
  if (shape.hasRank()) {
    for (int64_t dim : shape.getShape()) {
      proto->add_dim()->set_size(mlir::ShapedType::isDynamic(dim) ? -1 : dim);
    }
  } else {
    proto->set_unknown_rank(true);
  }
}

  QuantizationUnits GetQuantizableOps(ConstantOp op) const {
    // Non-float tensors do not need quantization.
    QuantizationUnits quantizable_ops;
    const ShapedType type = mlir::dyn_cast<ShapedType>(op.getType());
    if (!type || !type.getElementType().isF32()) return quantizable_ops;

    const Value value = op.getResult();

    for (OpOperand& use : value.getUses()) {

    return op.emitOpError("requires scalar output");

  return success();
}

OpFoldResult SizeOp::fold(FoldAdaptor) {
  ShapedType output_type = getType().cast<ShapedType>();
  if (!output_type.hasRank()) return {};
  ShapedType input_type = getOperand().getType().cast<ShapedType>();
  if (!input_type.hasStaticShape()) return {};
  int size = input_type.getNumElements();
  return DenseElementsAttr::get(

#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/core/errors.h"

namespace tensorflow {

using mlir::Builder;
using mlir::ShapedType;
using mlir::Type;

Status ConvertDataType(DataType dtype, Builder builder, Type* type) {
  switch (dtype) {
    case DT_HALF:
      *type = builder.getF16Type();
      return absl::OkStatus();

        element_type) {
      return op->emitOpError(
          "requires compatible element types for all operands and results");
    }
  }
  return success();
}

inline ShapedType MergeType(ShapedType a, ShapedType b) {
  if (!a.hasRank()) {
    return b;
  }
  if (!b.hasRank()) {
    return a;
  }
  int64_t rank = a.getRank();
  SmallVector<int64_t, 4> dims;
  dims.resize(rank);