import tensorflow  # pylint: disable=unused-import

from tensorflow.compiler.mlir.quantization.common.python import testing
from tensorflow.compiler.mlir.quantization.stablehlo import quantization_config_pb2 as stablehlo_quant_config_pb2
from tensorflow.compiler.mlir.quantization.tensorflow import quantization_options_pb2 as quant_opts_pb2
from tensorflow.compiler.mlir.quantization.tensorflow.python import quantize_model

include "mlir/Interfaces/SideEffectInterfaces.td"
include "tensorflow/compiler/mlir/lite/ir/tfl_op_interfaces.td"
include "tensorflow/compiler/mlir/lite/ir/tfl_op_enums.td"
include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td"

//===----------------------------------------------------------------------===//

#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"
#include "tensorflow/compiler/mlir/quantization/common/uniform_quantized_types.h"

#define DEBUG_TYPE "uniform-quantized-stablehlo-to-tfl"

namespace mlir {
namespace odml {
namespace {

// RUN:     --split-input-file --verify-diagnostics %s | FileCheck %s

// ============================================================================
// The following functions tests example quantization patterns outputted from
// JAX Quantizer. JAX Quantizer should output integer types, which are
// composed into `UniformQuantized{|PerAxis}Type` via
// `compose_uniform_quantized_type_pass.cc`.

#include "tensorflow/compiler/mlir/lite/utils/convert_type.h"
#include "tensorflow/compiler/mlir/lite/utils/utils.h"
#include "tensorflow/compiler/mlir/lite/utils/validators.h"
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/dynamic_shape_utils.h"

namespace mlir {
namespace TFL {

#include "tensorflow/compiler/mlir/lite/flatbuffer_operator.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/metrics/error_collector_inst.h"
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/lite/schema/mutable/schema_generated.h"
#include "tensorflow/compiler/mlir/lite/utils/convert_type.h"
#include "tensorflow/compiler/mlir/lite/utils/low_bit_utils.h"

  //
  // XlaCallModuleLoader uses mlir::stablehlo::deserializePortableArtifact,
  // which runs VhloLegalizeToStablehloPass whose depends on dialects:
  //   Func, Stablehlo, Shape, and Quantization.
  //
  // If we do not register them here, an error will be
  // triggered because we cannot load a dialect while in a
  // multi-threaded execution context, and PassManager is
  // multi-threaded.

    }
  }

  // Verify the quantized axis if the type is UniformQuantizedPerAxisType. Other
  // verifications to make sure the input and output has the same quantization
  // type, scale and zero point are performed by the SameOperandsAndResultsScale
  // trait.
  auto in_per_axis_qtype =
      QuantizedType::getQuantizedElementType(input_type)

func.func @testConcatInvalidScales(%arg0: tensor<*x!quant.uniform<i8:f32, 1.0>>, %arg1: tensor<*x!quant.uniform<i8:f32, 2.0>>) -> tensor<*x!quant.uniform<i8:f32, 1.0>> {
  // expected-error @+1 {{'tfl.concatenation' op quantization parameters violate the same scale constraint: !quant.uniform<i8:f32, 1.000000e+00> vs. !quant.uniform<i8:f32, 2.000000e+00>}}