#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_config.h"
#include "xla/service/hlo.pb.h"
#include "xla/service/hlo_parser.h"
#include "xla/translate/hlo_to_mhlo/hlo_to_mlir_hlo.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/graph_debug_info.pb.h"
#include "tensorflow/core/framework/types.pb.h"

#define TENSORFLOW_COMPILER_AOT_COMPILE_H_

#include <memory>
#include <string>

#include "tensorflow/compiler/aot/flags.h"
#include "tensorflow/compiler/tf2xla/tf2xla.pb.h"
#include "xla/service/cpu/cpu_compiler.h"
#include "xla/xla_data.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/platform/status.h"

namespace tensorflow {
namespace tfcompile {

#ifndef TENSORFLOW_COMPILER_JIT_RESOURCE_OPERATION_SAFETY_ANALYSIS_H_
#define TENSORFLOW_COMPILER_JIT_RESOURCE_OPERATION_SAFETY_ANALYSIS_H_

#include "xla/service/graphcycles/graphcycles.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/graph/graph.h"

namespace tensorflow {
// An XLA cluster hoists all resource reads to be beginning of the cluster
// execution and all the resource writes to the end.  This means it cannot

        "@llvm-project//mlir:FuncDialect",
        "@llvm-project//mlir:IR",
        "@local_tsl//tsl/platform:errors",
        "@local_xla//xla:shape_util",
        "@local_xla//xla/mlir_hlo:hlo_dialect_registration",
        "@local_xla//xla/translate/mhlo_to_hlo:type_to_shape",
    ],

#include "tensorflow/compiler/mlir/tf2xla/api/v2/tf_executor_to_graph.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/client_library.h"
#include "xla/client/compile_only_client.h"
#include "xla/stream_executor/host/host_platform_id.h"
#include "xla/stream_executor/platform_manager.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/framework/function.h"

}

Status XlaTensor::AllocateShapedBuffer(DataType dtype,
                                       const xla::Shape& on_device_shape,
                                       xla::LocalClient* client,
                                       int device_ordinal) {
  xla::Shape on_host_shape =
      xla::ShapeUtil::DeviceShapeToHostShape(on_device_shape);
  xla::ScopedShapedBuffer shaped_buffer(on_host_shape, on_device_shape,

    const std::vector<XlaCompiler::Argument>& args,
    DeviceCompileMode compile_mode, bool has_ref_vars,
    bool may_alias_resource_update,
    const XlaCompiler::CompilationResult** compilation_result,
    xla::PjRtClient** client, xla::PjRtLoadedExecutable** executable);

// Similar to the above function but it does not take a OpKernelContext.
// Instead, it takes the following arguments that are obtained from

#include "tensorflow/compiler/mlir/tensorflow/utils/convert_type.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/serialize_mlir_module_utils.h"
#include "xla/mlir_hlo/mhlo/IR/register.h"
#include "xla/shape.h"
#include "xla/translate/mhlo_to_hlo/type_to_shape.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tsl/platform/errors.h"

#ifndef TENSORFLOW_COMPILER_JIT_XLA_HOST_RECV_DEVICE_CONTEXT_H_
#define TENSORFLOW_COMPILER_JIT_XLA_HOST_RECV_DEVICE_CONTEXT_H_

#include "xla/shape.h"
#include "xla/stream_executor/device_memory.h"
#include "xla/stream_executor/stream.h"
#include "tensorflow/core/framework/device_base.h"
#include "tfrt/concurrency/async_value_ref.h"  // from @tf_runtime

namespace tensorflow {

// RUN: tf-quant-opt %s -split-input-file -quant-preprocess-op='target-opset=XLA quantization-method=weight_only enable-per-channel-quantization=false' | FileCheck --check-prefix PerTensor %s
// RUN: tf-quant-opt %s -split-input-file -quant-preprocess-op='target-opset=XLA quantization-method=weight_only enable-per-channel-quantization=true' | FileCheck --check-prefix PerChannel %s

module {
  // For XLA weight-only per-channel depthwise convolution, tensor shape should have