// method in textproto format. The dot_general op quantization is explicitly
// disabled by having `_quantization_method = "no_quantization { }"`.
// DISABLE-ALL-DOT-GENERAL-SAME: _entry_function = @composite_dot_general_fn_1
// DISABLE-ALL-DOT-GENERAL-SAME: _original_entry_function
// DISABLE-ALL-DOT-GENERAL-SAME: _quantization_method = "no_quantization { }"
// DISABLE-ALL-DOT-GENERAL-SAME: _tfl_quant_trait = "fully_quantizable"

}

TEST_F(LiftAsFunctionCallTest, GetQuantizationMethodSucceeds) {
  // Function containing a simple `TF::XlaCallModuleOp` with a valid string
  // attribute `_quantization_method` set to `"no_quantization {}"`.
  constexpr absl::string_view kXlaCallModuleOpWithQuantizationMethodAttr =
      R"mlir(
    func.func @main(%arg0: tensor<1x1x3xf32>, %arg1: tensor<3x4xf32>) -> tensor<1x1x4xf32> {

// quantizable units.
constexpr absl::string_view kSpecsDisableAllDotGeneral =
    R"pb(specs
         [ {
           matcher { function_name { regex: "composite_dot_general_.*" } }
           method { no_quantization {} }
         }])pb";

// Configure `QuantizationSpecs` to apply `StaticRangePtq` to all quantizable
// units.
constexpr absl::string_view kSpecsStaticRangePtqToAll =
    R"pb(specs
         [ {

  EXPECT_TRUE(IsOpQuantizableStableHlo(xla_call_module_op));
}

TEST_F(IsOpQuantizableStableHloTest, DenylistedXlaCallModuleOpNotQuantizable) {
  // A `TF::XlaCallModuleOp` with `_quantization_method = "no_quantization {}"`,
  // indicating it has been explicitly denylisted by the user.
  constexpr absl::string_view kModuleDenylistedXlaCallModuleOp = R"mlir(

  //   matcher {function_name {regex: "composite_conv.*"}}
  //   method {static_range_ptq {...}}}
  // }
  // specs {
  //   matcher {function_name {regex: "composite_dot_general_fn_1"}}
  //   method {no_quantization {}}
  // }
  const QuantizationConfig new_config = ExpandPresets(config);
  ASSERT_THAT(new_config.specs().specs(), SizeIs(3));

  const QuantizationSpec& first_spec = new_config.specs().specs(0);

                        function_name=qc.FunctionNameMatcherSpec(
                            regex='composite_dot_general.*'
                        )
                    ),
                    method=qc.Method(no_quantization={}),
                )
            ]
        ),
    )
    quantization.quantize_saved_model(
        self._input_saved_model_path,
        self._output_saved_model_path,
        config,
    )

  // Matches lifted functions by their names.
  FunctionNameMatcherSpec function_name = 1;
}

// Specifies how to quantize matched quantizable units.
message Method {
  oneof method {
    NoQuantization no_quantization = 1;
    StaticRangePtq static_range_ptq = 2;
    WeightOnlyPtq weight_only_ptq = 3;
  }
}

// A QuantizationSpec is essentially a (matcher spec, quantization method) pair,

using ::stablehlo::quantization::Method;
using ::stablehlo::quantization::StaticRangePtq;

// Whether it represents a lifted function (i.e. `op` is the corresponding
// `XlaCallModuleOp`) that is explicitly marked `NoQuantization`.
bool IsDenylistedLiftedFunction(Operation* op) {
  if (auto xla_call_module_op = dyn_cast_or_null<TF::XlaCallModuleOp>(op);
      xla_call_module_op != nullptr) {

limitations under the License.
==============================================================================*/

// This is the quantization definition file for TensorFlow.

#ifdef TF_Quantization
#else
#define TF_Quantization

include "mlir/IR/OpBase.td"
include "mlir/Dialect/Quant/QuantOpsBase.td"

//===----------------------------------------------------------------------===//
// QuantizedType definitions.

// `xla_call_module_op` is a quantizable unit, then a `QuantizationResult` will
// be returned with its `name` field set to the callee's name. The `method`
// field will be set to `NoQuantization` because remaining `xla_call_module_op`s
// means they are not quantized. Returns `std::nullopt` if `xla_call_module_op`
// is not a quantizable unit.
std::optional<QuantizationResult> GetQuantizationResult(