#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/client/compile_only_client.h"
#include "xla/hlo/ir/hlo_input_output_alias_config.h"
#include "xla/mlir_hlo/mhlo/IR/register.h"
#include "xla/shape.h"
#include "xla/shape_util.h"
#include "xla/status_macros.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/op.h"

            step_marker_location = ""} : (tensor<*xi32>) -> tensor<*xi32>
        return
      }
      func @func(%arg0: tensor<*xi32> {mhlo.sharding = "\01\02\03"}) ->
                (tensor<*xi32> {mhlo.sharding = "\08\01\1A\01\01\22\01\00"}) {
        %0 = "tf.XlaSharding"(%arg0) {_XlaSharding = "\01\02\03", sharding = "\01\02\03"} : (tensor<*xi32>) -> tensor<*xi32>

// operations constraints to form a valid cluster.
//
// Example: compilation using XLA (MHLO) lowering
//
//   %0 = "tf.Transpose"(%input, %perm)
//        : (tensor<?x?xf32>, tensor<2xi32>) -> tensor<?x?xf32>
//
//   XLAs `mhlo.transpose` operation requires permutation to be an attribute
//   (compile time value), so it means that if we want to put `tf.Transpose`

#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "xla/mlir_hlo/mhlo/transforms/passes.h"

namespace tensorflow {
namespace quantization {
namespace {

using ::tensorflow::quantization::QuantizationOptions;

void AddConvertTpuToCpuModelPasses(mlir::OpPassManager &pm) {

  let dependentDialects = ["mlir::stablehlo::StablehloDialect"];
  let constructor = "mlir::quant::stablehlo::CreateQuantizeWeightPass()";
}

def UnfuseMhloBatchNormPass : Pass<"stablehlo-unfuse-mhlo-batch-norm", "mlir::func::FuncOp"> {
  let summary = "Unfuses batch normalization into arithmetic ops.";
}

def LiftQuantizableSpotsAsFunctionsPass : Pass<"stablehlo-lift-quantizable-spots-as-functions", "mlir::ModuleOp"> {

#include "llvm/ADT/DenseSet.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tpu_embedding_ops_registry.h"

namespace mlir {
namespace mhlo {

namespace {

// Returns ops that should use MLIR legalization.
// All other ops not in this list should use XlaOpKernel.
const llvm::DenseSet<mlir::TypeID>& MlirAlwaysOps() {

  if (HasStringResult(op)) return false;
  if (MatchesPattern(op, supported_ops)) return true;

  auto abstractOp = op.getRegisteredInfo();
  if (!abstractOp) return false;
  return mlir::mhlo::HasTf2XlaFallback(abstractOp->getTypeID());
}

bool IsVariant(Value value) {
  return mlir::isa<mlir::TF::VariantType>(
      getElementTypeOrSelf(value.getType()));
}

                      op.getNumOperands(), input_shardings.size()));

  // Set args metadata in proto.
  mlir::StringAttr replication_attr_name = mlir::StringAttr::get(
      op.getContext(), "mhlo.is_same_data_across_replicas");

  auto dynamic_arg_idx = op->getAttrOfType<ArrayAttr>(TF::kDynamicArgIndexAttr);
  llvm::SmallSet<int, 4> dynamic_arg_idx_set;
  if (dynamic_arg_idx) {

namespace mlir {
namespace quant {
class QuantizationDialect;
}
namespace quantfork {
class QuantizationForkDialect;
}
namespace mhlo {
class MhloDialect;
}
namespace TF {
class TensorFlowDialect;
}
namespace TFL {
class TFLDialect;
typedef TFLDialect TensorFlowLiteDialect;
}  // namespace TFL
namespace func {
class FuncOp;
}

#include "tensorflow/compiler/mlir/tensorflow/utils/xla_call_module_attrs.h"
#include "tensorflow/compiler/tf2xla/kernels/xla_call_module_loader.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"  // IWYU pragma: keep
#include "tsl/platform/statusor.h"

namespace mlir {
namespace TF {
namespace {

#define GEN_PASS_DEF_XLACALLMODULEDESERIALIZATIONPASS