}

  MLIRContext ctx_;
};

TEST_F(CreateI8F32UniformQuantizedPerAxisTypeTest, I8StorageTypeSucceeds) {
  const UniformQuantizedPerAxisType quantized_type =
      CreateI8F32UniformQuantizedPerAxisType(
          UnknownLoc::get(&ctx_), ctx_,
          /*scales=*/SmallVector<double, 2>{1.0, 1.0},
          /*zero_points=*/SmallVector<int64_t, 2>{0, 0},
          /*quantization_dimension=*/0);

// zero point values can be non-zero values.
// If `narrow_range` is set true (ex: for weights), a restricted range of
// integers will be used for symmetric mapping, i.e. [-127, 127].
UniformQuantizedPerAxisType CreateI8F32UniformQuantizedPerAxisType(
    Location loc, MLIRContext& context, ArrayRef<double> scales,
    ArrayRef<int64_t> zero_points, int quantization_dimension,
    bool narrow_range = false);

      /*expressedType=*/FloatType::getF32(&context), scale, zero_point,
      /*storageTypeMin=*/llvm::minIntN(32),
      /*storageTypeMax=*/llvm::maxIntN(32));
}

UniformQuantizedPerAxisType CreateI8F32UniformQuantizedPerAxisType(
    const Location loc, MLIRContext& context, const ArrayRef<double> scales,
    const ArrayRef<int64_t> zero_points, const int quantization_dimension,
    const bool narrow_range) {

#define DEBUG_TYPE "stablehlo-compose-uniform-quantized-type"

namespace mlir {
namespace odml {
namespace {

using ::mlir::quant::CreateI8F32UniformQuantizedPerAxisType;
using ::mlir::quant::CreateI8F32UniformQuantizedType;
using ::mlir::quant::TryCast;
using ::mlir::quant::UniformQuantizedPerAxisType;
using ::mlir::quant::UniformQuantizedType;

using ::mlir::quant::CastI64ArrayToI32;
using ::mlir::quant::CastI64ToI32;
using ::mlir::quant::CreateI32F32UniformQuantizedPerAxisType;
using ::mlir::quant::CreateI32F32UniformQuantizedType;
using ::mlir::quant::CreateI8F32UniformQuantizedPerAxisType;
using ::mlir::quant::CreateI8F32UniformQuantizedType;
using ::mlir::quant::FindOperandOfType;
using ::mlir::quant::FindUserOfType;
using ::mlir::quant::GetElementType;

          "new quantization dimension not found in perm");
    }

    // Insert a QDQ pair with the new quantized dimension after TransposeOp
    auto new_qtype = quant::CreateI8F32UniformQuantizedPerAxisType(
        transpose_op.getLoc(), *rewriter.getContext(), aqtype.getScales(),
        aqtype.getZeroPoints(), new_out_quant_dim, /*narrow_range=*/true);
    auto new_tensor_type = RankedTensorType::getChecked(