//===----------------------------------------------------------------------===//
// XlaCallModule
//===----------------------------------------------------------------------===//

// Creates a pass that deserializes functions in the StableHLO modules from
// `tf.XlaCallModule` to the top-level module.
std::unique_ptr<OperationPass<ModuleOp>>
CreateXlaCallModuleDeserializationPass();

    // MIN-MAX-CHECK: %[[ARG1_AGG:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%[[ARG1_ID]])
    // MIN-MAX-CHECK-SAME: id = "composite_dot_general_fn_1_arg_1_calibration_method_1"
    // MIN-MAX-CHECK: %[[RES:.*]] = "tf.XlaCallModule"(%[[ARG0_AGG]], %[[ARG1_AGG]])
    // MIN-MAX-CHECK: %[[RES_AGG:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%[[RES]])
    // MIN-MAX-CHECK-SAME: id = "composite_dot_general_fn_1_calibration_method_1"

namespace mlir {
namespace odml {

static constexpr std::string_view kStablehloModuleDefaultEntryFuncName = "main";
static constexpr std::string_view kStablehloFuncNamePrefix = "XlaCallModule";
static constexpr char kShardingAttr[] = "mhlo.sharding";
static constexpr char kShardingName[] = "Sharding";

class RemoveCustomCallWithSharding
    : public OpRewritePattern<stablehlo::CustomCallOp> {

      %0 = stablehlo.constant dense<[-0.211145893, -0.708605706]> : tensor<2xf32>
      %1 = stablehlo.constant dense<[[-0.630731344, 0.54962182], [0.180364341, -0.764542698]]> : tensor<2x2xf32>

    %0 = "tf.XlaCallModule"(%output, %cst) <{Sout = [#tf_type.shape<10x1x3>], dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64}> {_entry_function = @composite_dot_general_with_relu_fn_1,...

// Default platform for XlaCallModuleOp.
constexpr StringRef kPlatformCpu = "CPU";
// Name of `tf.XlaCallModule`'s dictionary attribute for keeping the
// deserialized stablehlo module's attributes.
constexpr StringRef kStablehloModuleAttrsAttrName = "_stablehlo_module_attrs";
// Attribute required for running shape refinement pass enabled in XlaCallModule
// version 8 and above.
constexpr StringRef kUsesShapePolymorphismAttr = "jax.uses_shape_polymorphism";

    // legalizing to stablehlo.constant.
    AddTFToStablehloPasses(pm_after_freezing_variables, input_arg_shapes);
  }

  if (deserialize_xla_call_module) {
    // Deserialize the StableHLO module embedded in tf.XlaCallModule and lifts
    // the StableHLO functions to the top level module. This is needed for
    // StableHLO quantization. Also restores some shape information for

  constexpr absl::string_view kNonQuantizedDotGeneral = R"mlir(
    func.func @main(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> {
      %0 = stablehlo.constant dense<3.000000e+0> : tensor<2x3xf32>

inline FlatSymbolRefAttr GetFuncAttr(TF::XlaCallModuleOp call_op) {
  return call_op->getAttrOfType<FlatSymbolRefAttr>(
      TF::kStablehloEntryFunctionAttrName);
}

// Returns the entry function name for the given tf.XlaCallModule op. Returns
// empty string if such attribute does not exist.
StringRef GetEntryFunctionName(TF::XlaCallModuleOp op);

// Checks whether the given op contains QuantizationTrait::FullyQuantizable.

     function block to be compatible with SavedModel structure.
  }];
}

def RestoreFunctionNamePass : Pass<"stablehlo-restore-function-name", "ModuleOp"> {
  let summary = "Restores function name from XlaCallModule op.";
}

def QuantizeCompositeFunctionsPass : Pass<"stablehlo-quantize-composite-functions", "ModuleOp"> {
  let summary = "Quantize composite functions with QDQ input / outputs.";
  let options = [