// RUN: tf-opt --tfl-decompose-hybrid-quantization --verify-each %s | FileCheck %s

// CHECK-LABEL: @test_conv2d_float
func.func @test_conv2d_float(%arg0: tensor<1x32x32x8xf32>) -> tensor<1x32x32x16xf32> {
  // CHECK-DAG: %[[VAL0:.+]] = "tfl.pseudo_const"() <{value = dense<42> : tensor<16x1x1x8xi8>}>
  // CHECK-DAG: %[[VAL1:.+]] = "tfl.pseudo_const"() <{value = dense<1> : tensor<16x1x1x8xi8>}>

    }

    def "can apply hybrid plugin with no id"() {
        when:
        manager.apply(hybridClass)

        then:
        1 * target.applyImperativeRulesHybrid(null, { hybridClass.isInstance(it) }, hybridClass)

        and:
        manager.pluginContainer.size() == 1
    }

    def "can apply hybrid plugin by class with id"() {
        given:

      attrs.push_back(rewriter.getNamedAttr(
          attr_minmax, rewriter.getI64IntegerAttr(quant_val)));
    }
  }
  return success();
}

// This LogicalResult covers both the hybrid and fully quantized op cases.
LogicalResult FillAttributesForUniformQuantizedDotOp(
    PatternRewriter& rewriter, Operation* op,
    llvm::StringMap<Attribute>& identifier_to_attr,

   (IsConstTensor $filter),
   (IsInt32ElementType $conv),
   (HasStaticShapeConstraint $filter),
   (HasStaticShapeAtDimsConstraint<"3"> $input)],
  [], (addBenefit 10)>;

// Convert Conv2D with hybrid inputs (f32 activation/int8 weight) to XlaConv
def ConvertTFConv2DToXLAConvOpWeightOnly : Pat<
  (TF_Conv2DOp:$conv
    $input,
    (TF_MulOp (TF_CastOp (TF_IdentityOp $filter), $truncate1), $scale),

		// cluster name, generate the jwt filter config using remote Jwks.
		// If failed to parse the cluster name, only fallback to let istiod to fetch the jwksUri when
		// remoteJwksMode is Hybrid.
		if features.JwksFetchMode != jwt.Istiod && jwtRule.JwksUri != "" {
			jwksInfo, err := security.ParseJwksURI(jwtRule.JwksUri)
			if err != nil {
				authnLog.Errorf("Failed to parse jwt rule jwks uri %v", err)
			}

		v := env.Register(
			"PILOT_JWT_ENABLE_REMOTE_JWKS",
			"false",
			"Mode of fetching JWKs from JwksUri in RequestAuthentication. Supported value: "+
				"istiod, false, hybrid, true, envoy. The client fetching JWKs is as following: "+
				"istiod/false - Istiod; hybrid/true - Envoy and fallback to Istiod if JWKs server is external; "+
				"envoy - Envoy.",
		).Get()
		return jwt.ConvertToJwksFetchMode(v)
	}()

      double error_tolerance = 0.1) {
    // Expected result is calculated by evaluating using TF kernels. In some
    // cases, TF kernel behaves differently from lowered graph (e.g. Hybrid
    // ops). So we optionally use a different graph to calculate the expected
    // result.
    TF_ASSERT_OK_AND_ASSIGN(
        auto expected,
        this->EvaluateTfFunction(

  using impl::InsertWeightParamPassBase<
      InsertWeightParamPass>::InsertWeightParamPassBase;

 private:
  void runOnOperation() override;
};

// Inserts quantization parameters for weights for hybrid quantization of
// `stablehlo.convolution` and `stablehlo.dot_general`.
class InsertWeightParamPattern
    : public OpTraitRewritePattern<OpTrait::ConstantLike> {
 public:

  ];
}
def DecomposeHybridQuantizationPass : Pass<"tfl-decompose-hybrid-quantization", "mlir::func::FuncOp"> {
  let summary = "Decomposes hybridge quantization to explicit quantize / dequantize";
  let description = [{
      Decomposes (with explicit quantize/dequantize ops) selected math
      operations which exist in the model with hybrid quantization
      (some arguments/results left in floating point).
  }];

// Creates an instance of the TensorFlow Lite dialect QuantizeVariables pass.
std::unique_ptr<OperationPass<ModuleOp>> CreatePrepareQuantizeVariablesPass();

// Creates an instance of the TensorFlow Lite pass that decomposes hybrid
// quantization patterns to the same dense operation with tfl dequantization
// and quantization patterns.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateDecomposeHybridQuantizationPass();