#include "tensorflow/compiler/jit/variable_info.h"
#include "tensorflow/compiler/jit/xla_tensor.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/local_client.h"
#include "xla/pjrt/pjrt_client.h"
#include "xla/service/shaped_buffer.h"
#include "xla/stream_executor/device_memory_allocator.h"
#include "tensorflow/core/framework/allocation_description.pb.h"
#include "tensorflow/core/framework/tensor.h"

#include "tensorflow/compiler/mlir/op_or_arg_name_mapper.h"
#include "tensorflow/compiler/tf2xla/xla_context.h"
#include "tensorflow/compiler/tf2xla/xla_expression.h"
#include "xla/client/xla_builder.h"
#include "xla/client/xla_computation.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/framework/op_kernel.h"

namespace mlir {
namespace mhlo {

#include "tensorflow/compiler/jit/xla_compiler_options_util.h"

#include "xla/pjrt/pjrt_client.h"
#include "tensorflow/core/framework/function.h"
#include "tsl/framework/device_id_utils.h"

namespace tensorflow {
namespace {
using XlaDeviceCompiler =
    DeviceCompiler<xla::LocalExecutable, xla::LocalClient>;
using PjRtDeviceCompiler =
    DeviceCompiler<xla::PjRtLoadedExecutable, xla::PjRtClient>;

#include "tensorflow/compiler/tf2xla/const_analysis.h"
#include "tensorflow/compiler/tf2xla/resource_operation_table.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/service/graphcycles/graphcycles.h"
#include "xla/statusor.h"
#include "xla/union_find.h"
#include "xla/util.h"
#include "tensorflow/core/common_runtime/function.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/framework/attr_value.pb.h"

      $resource,
      (TF_SubOp
        (CreateTFReadVariableOp $src_op, $value, $resource),
        $value
      ),
      (CreateConstBoolAttrFalse)
    )
  >;

// This decomposition is only correct inside XLA as it ignores use_locking
// attribute.
def DecomposeResourceApplyGradientDescentOp :
  Pat<
    (TF_ResourceApplyGradientDescentOp:$src_op $resource, $alpha, $delta, BoolAttr:$_),
    (TF_AssignVariableOp
      $resource,

// The XlaDevice executes a TensorFlow graph using the XLA linear algebra
// runtime.
//
// Operators assigned to an XlaDevice are compiled into XLA computations.
// Tensors on an XlaDevice are thin wrappers around XLA ScopedShapedBuffers.
//
// XlaDevice is instantiated separately for each XLA backend (e.g., CPU or GPU),
// under different names (e.g., XLA_CPU or XLA_GPU).

  EXPECT_EQ(exec_args.size(), 2);

  std::shared_ptr<xla::Literal> literal1 = *exec_args[0]->ToLiteralSync();
  EXPECT_TRUE(xla::LiteralTestUtil::Equal(
      *literal1, xla::LiteralUtil::CreateR2<int32_t>({{1, 2, 3}})));

  std::shared_ptr<xla::Literal> literal2 = *exec_args[1]->ToLiteralSync();
  EXPECT_TRUE(xla::LiteralTestUtil::Equal(
      *literal2, xla::LiteralUtil::CreateR2<int32_t>({{4, 5, 6}})));
}

#include "tensorflow/compiler/jit/xla_platform_info.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/executable_build_options.h"
#include "xla/client/local_client.h"
#include "xla/service/hlo_graph_dumper.h"
#include "xla/status_macros.h"
#include "xla/stream_executor/host/host_platform_id.h"
#include "xla/stream_executor/platform.h"
#include "tensorflow/core/common_runtime/eager/tensor_handle.h"

//     eliminates the use of resource variables.
//   . Legalizes the operations to XLA HLO operations.
//   . Canonicalizes the XLA HLO operations.
//
// device_type: XLA JIT device to use for compilation such as "XLA_CPU_JIT",
//   "XLA_GPU_JIT" or "XLA_TPU_JIT".
// use_tuple_args: when this is true, always create a tuple argument for the
//   entry computation.

        "@local_xla//xla/client:xla_computation",
        "@local_xla//xla/hlo/ir:hlo",
        "@local_xla//xla/mlir_hlo",
        "@local_xla//xla/mlir_hlo:hlo_dialect_registration",
        "@local_xla//xla/mlir_hlo:mhlo_passes",
        "@local_xla//xla/translate/mhlo_to_hlo:layout_util",
        "@local_xla//xla/translate/mhlo_to_hlo:mlir_hlo_to_hlo",
        "@local_xla//xla/translate/mhlo_to_hlo:type_to_shape",