std::unique_ptr<OperationPass<func::FuncOp>> CreateImportQuantStatsPass(
    OperationToName op_to_name, const std::string& stats_str);

// Creates an instance pass to import quantization stats to the operations in
// the function. A custom method to get the name from the op is used because
// different dialect ops might have different ways to assign the name.
std::unique_ptr<OperationPass<func::FuncOp>>

    cluster.walk([](Operation* op) { DropWhileShapeInvariantAttr(op); });
  });
}
}  // namespace

std::unique_ptr<OperationPass<func::FuncOp>>
CreateDropWhileShapeInvariantPass() {
  return std::make_unique<DropWhileShapeInvariantPass>();
}

std::unique_ptr<OperationPass<func::FuncOp>>
CreateDropWhileShapeInvariantInDeviceClusterPass() {

// also owns a PjRtBuffer for XLA kernel's usage.
class PjRtTensorBuffer : public TensorBuffer {
 public:
  PjRtTensorBuffer(const void* ptr, size_t expected_size,
                   std::unique_ptr<xla::PjRtBuffer> pjrt_buffer)
      : TensorBuffer(const_cast<void*>(ptr)),
        expected_size_(expected_size),
        pjrt_buffer_(std::move(pjrt_buffer)) {}

  size_t size() const override { return expected_size_; }

  return TensorHandlePtr(result_handle);
}

template <typename value_type>
void ExpectScalarEq(TFE_TensorHandle* handle, value_type expected_value) {
  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
      TF_NewStatus(), TF_DeleteStatus);
  std::unique_ptr<TF_Tensor, decltype(&TF_DeleteTensor)> actual_value(
      TFE_TensorHandleResolve(handle, status.get()), TF_DeleteTensor);

#include "tensorflow/core/tfrt/common/async_value_tensor.h"
#include "tensorflow/core/tfrt/common/create_pjrt_client_util.h"
#include "tsl/framework/device_id_utils.h"

namespace tensorflow {
namespace {

absl::StatusOr<std::unique_ptr<xla::PjRtBuffer>> HostTensorToPjRtBuffer(
    const tensorflow::Tensor* cpu_tensor, tensorflow::Device* device,
    xla::PjRtClient* pjrt_client,
    const XlaShapeLayoutHelpers::ShapeDeterminationFns

class FuncOp;
}  // namespace func

namespace quantfork {

/// Creates a pass that converts quantization simulation operations (i.e.
/// FakeQuant and those like it) to casts into/out of supported QuantizedTypes.
std::unique_ptr<OperationPass<func::FuncOp>> createConvertSimulatedQuantPass();

/// Creates a pass that converts constants followed by a qbarrier to a
/// constant whose value is quantized. This is typically one of the last

class XlaDeviceCompilerClient
    : public DeviceCompilerClient<xla::LocalExecutable, xla::LocalClient> {
 public:
  explicit XlaDeviceCompilerClient(xla::LocalClient* client)
      : client_(client) {}

  absl::StatusOr<std::unique_ptr<xla::LocalExecutable>> BuildExecutable(
      const XlaCompiler::Options& options,
      const XlaCompiler::CompilationResult& result) override;

  // Returns a serialized AOT result obtained by exporting the available

// Returns the spec for the given operation that can be used for both of
// dynamic and static range quantization.
std::unique_ptr<OpQuantSpec> GetTFOpQuantSpec(Operation* op);

// Returns quantization scale specs (fixed output, same scale) for a TF op.
std::unique_ptr<OpQuantScaleSpec> GetTfQuantScaleSpec(Operation* op);

}  // namespace quant
}  // namespace mlir

Status LoadTFConcreteFunction(
    const SavedConcreteFunction& saved_concrete_function,
    const FunctionDef* function_def,
    const std::unordered_map<int, std::unique_ptr<TensorHandleConvertible>>&
        captured_objects,
    ImmediateExecutionContext* ctx, std::unique_ptr<TFConcreteFunction>* out) {
  TF_RETURN_IF_ERROR(ValidateSavedFunctionCompatibleWithFunctionDef(
      saved_concrete_function, function_def));

Status TestTensorHandleWithDims(AbstractContext* ctx, const T* data,
                                const int64_t* dims, int num_dims,
                                AbstractTensorHandle** tensor) {
  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
      TF_NewStatus(), TF_DeleteStatus);
  TFE_Context* eager_ctx =
      TF_ExecutionContextGetTFEContext(wrap(ctx), status.get());