#include "mlir/Support/LLVM.h"  // from @llvm-project

namespace mlir {
namespace TFL {

FloatAttr ExtractSingleElementAsFloat(ElementsAttr attr) {
  if (attr.getShapedType().getNumElements() != 1 ||
      !mlir::isa<FloatType>(attr.getShapedType().getElementType())) {
    return {};
  }
  return attr.getSplatValue<FloatAttr>();
}

FloatAttr GetSingleElementAsFloatOrSelf(Attribute attr) {

  auto tensorArg = mlir::dyn_cast<TensorType>(getArg().getType());
  if (!tensorArg) return emitOpError("arg needs to be tensor type.");

  // Verify layerStats attribute.
  {
    auto layerStatsType = getLayerStats().getShapedType();
    if (!mlir::isa<FloatType>(layerStatsType.getElementType())) {
      return emitOpError("layerStats must have a floating point element type");
    }

  auto tensorArg = mlir::dyn_cast<TensorType>(getArg().getType());
  if (!tensorArg) return emitOpError("arg needs to be tensor type.");

  // Verify layerStats attribute.
  {
    auto layerStatsType = getLayerStats().getShapedType();
    if (!mlir::isa<FloatType>(layerStatsType.getElementType())) {
      return emitOpError("layerStats must have a floating point element type");
    }

// Quantize attribute $0 by using quantization parameter from %1.
def QuantizeByQuantizedType : NativeCodeCall<"quant::Quantize($0, $1.getValue())">;
def F32ElementsAttr : ElementsAttrBase<
  CPred<"$_self.cast<ElementsAttr>().getShapedType().getElementType().isF32()">, "float constant tensor">;

// Squash tfl.dequantize and tfl.quantize pairs.
// TODO(fengliuai): Compare the scale of input and output. This can also be

    // Only fuses multiplier if all dimensions other than the out channel
    // dimension are equal to 1.
    if (!TFL::IsDimensionsDegenerateExceptLastOne(
            mul_value.getShapedType().getShape())) {
      return rewriter.notifyMatchFailure(mul_op, [&](::mlir::Diagnostic &diag) {
        diag << "entities 'mul_value' failed to satisfy constraint: "
                "unsupported dimensions";
      });

    assert(perm_tensor.getType().getRank() == 1);
    const int num_dimensions = input_tensor.getShapedType().getRank();
    assert(perm_tensor.getType().getNumElements() == num_dimensions);

    ArrayRef<int64_t> input_shape = input_tensor.getShapedType().getShape();
    auto output_type = mlir::cast<ShapedType>(op.getOutput().getType());

    SmallVector<int32_t, 4> perm;

def CanFuseMulAndConv2D :
    Constraint<CPred<"TFL::IsBroadcastableElementsAttrs($0, $1) && TFL::IsDimensionsDegenerateExceptLastOne($1)">>;

def F32ElementsAttr : ElementsAttrBase<
    CPred<"$_self.cast<ElementsAttr>().getShapedType().getElementType().isF32()">, "float constant tensor">;
def DefinedByConv2D : Constraint<CPred<"llvm::isa_and_nonnull<mlir::TF::Conv2DOp>($0.getDefiningOp())">>;
// Checks if the value has only one user.

    const auto value_attr =
        mlir::cast<DenseFPElementsAttr>(const_op.getValue());
    const ArrayRef<int64_t> original_shape =
        value_attr.getShapedType().getShape();

    const SmallVector<float> original_values =
        llvm::to_vector(value_attr.getValues<float>());

    // Fold the constant value by transposing the values according to the

    llvm::SmallVector<int64_t> original_dimensions;
    for (auto dim : input_tensor_shape) original_dimensions.push_back(dim.size);
    return ElementsAttr(mlir::SplatElementsAttr::get(
        single_attr.getShapedType().clone(original_dimensions),
        single_attr.getValues<mlir::Attribute>()[0]));
  }

  Tensor t;
  if (!t.FromProto(input_tensor))
    return InvalidArgument("Failed to parse input_tensor.");

    tensor.flat<T>().setValues(values);

    auto value_or = ConvertTensor(tensor, &b);
    TF_ASSERT_OK(value_or.status());
    auto attr = value_or.value();

    EXPECT_EQ(attr.getShapedType().getElementType(), expected_ty);

    Tensor out;
    TF_ASSERT_OK(ConvertToTensor(attr, &out));

    test::ExpectTensorEqual<T>(tensor, out);
  }
};

TEST_F(ConvertTensorTest, Simple) {