#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "tensorflow/core/platform/statusor.h"

namespace tensorflow {

class ProcessFunctionLibraryRuntime;
class Device;
class Tensor;
class TensorHandle;
class EagerContext;

enum class IrExportStage {
  HLO,
  HLO_NO_METADATA,
  HLO_SERIALIZED,
  OPTIMIZED_HLO,
  OPTIMIZED_HLO_SERIALIZED,
  OPTIMIZED_HLO_PROTO_SERIALIZED,
  OPTIMIZED_HLO_DOT
};

#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"

namespace tensorflow {
namespace {

ImmediateTensorHandlePtr CreateScalarTensorHandle(EagerContext* context,
                                                  float value) {
  AbstractTensorPtr tensor(context->CreateFloatScalar(value));
  ImmediateTensorHandlePtr handle(context->CreateLocalHandle(tensor.get()));

  } else {
    std::unique_ptr<tensorflow::TFSavedModelAPI> saved_model;
    status->status = tensorflow::TFSavedModelAPI::Load(
        dirname, absl::nullopt,
        tensorflow::down_cast<tensorflow::EagerContext*>(
            tensorflow::unwrap(ctx)),
        &saved_model);
    result = std::move(saved_model);
  }

  if (!status->status.ok()) {
    return nullptr;
  }
  return tensorflow::wrap(result.release());

        new Variable(ctx, dtype, shape, std::move(name), std::move(handle)));
    return Status();
  }

  if (!tensorflow::isa<EagerContext>(ctx)) {
    return errors::InvalidArgument(
        "Can only load distributed variables with EagerContext.");
  }

  EagerContext* eager_ctx = reinterpret_cast<EagerContext*>(ctx);

  std::vector<TensorHandle*> handles;
  for (const auto& device : component_devices) {

#include "tensorflow/core/platform/types.h"

namespace tensorflow {
namespace testing {

// Creates a DeviceMgr suitable for local tests.
std::unique_ptr<StaticDeviceMgr> CreateTestingDeviceMgr();

// Creates an EagerContext suitable for local tests. Does not take ownership
// of `device_mgr`.
EagerContextPtr CreateTestingEagerContext(DeviceMgr* device_mgr);

// Converts a tensorflow::DatatypeSet to std::vector<DataType>.

  }
  return absl::OkStatus();
}

absl::StatusOr<std::vector<XlaCompiler::Argument>> PrepareXlaCompilerArgs(
    FunctionLibraryRuntime* flr, const NameAttrList& function,
    EagerContext* context, Device* dev,
    absl::Span<const ArgShapeAndDType> input_arg_shape_and_dtype,
    absl::Span<const TensorHandle* const> input_handles,
    CompilerArgSource compiler_arg_source) {

  auto r = tsl::core::RefCountPtr<tensorflow::IntraProcessRendezvous>(
      new tensorflow::IntraProcessRendezvous(device_mgr.get()));
  tensorflow::EagerContext* eager_context = new tensorflow::EagerContext(
      opts->session_options.options,
      static_cast<tensorflow::ContextDevicePlacementPolicy>(
          opts->device_placement_policy),
      opts->async, device_mgr.release(),

  LOG_FIRST_N(INFO, 1) << "Run Node Expansion Passes";

  // Get the FunctionDef and insert that into the context
  const NodeDef& ndef = orig_op->MutableAttrs()->BuildNodeDef();
  auto& ctx = orig_op->EagerContext();
  Fprint128 cache_key =
      orig_op->MutableAttrs()->CacheKey(orig_op->DeviceName());
  // Include soft placement policy in cache key since the placement strategy
  // can change and thus affect which kernel is picked.

        options, "/job:localhost/replica:0/task:0", &devices));
    device_mgr_ = absl::make_unique<StaticDeviceMgr>(std::move(devices));
    ctx_ = testing::CreateTestingEagerContext(device_mgr_.get());
  }

  EagerContext* context() { return ctx_.get(); }

 private:
  std::unique_ptr<StaticDeviceMgr> device_mgr_;
  EagerContextPtr ctx_;
};

// Sanity check that constructing a tensorflow::Variable from a SavedVariable

// FlatTensorFunction models a TF2 eager runtime view of a callable function,
// taking + returning flat lists of tensors, including any captures.
// Effectively, it is a thin wrapper around a FunctionDef owned by the
// EagerContext, and any TensorHandle captures associated with the function. The
// MakeCallOp method handles the logic of marshaling captures after the user
// provided inputs automatically.