#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"

namespace mlir {
namespace TF {
namespace {

constexpr char kDeviceAttr[] = "device";
constexpr char kTFDeviceAttr[] = "tf.device";

#define GEN_PASS_DEF_SIMPLETFDEVICEASSIGNMENTPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

class SimpleTFDeviceAssignmentPass

#include "tensorflow/compiler/mlir/tensorflow/utils/convert_tensor.h"

namespace mlir {
namespace TF {

namespace {

constexpr const char *kDeviceAttr = "device";
constexpr const char *kTFDeviceAttr = "tf.device";

#define GEN_PASS_DEF_TENSORDEVICECOPYCONVERSIONPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

struct TensorDeviceCopyConversionPass

    // function.
    for (int i : llvm::seq<int>(0, metadata.input_devices.size())) {
      func_op.setArgAttr(i, kTFDeviceAttr,
                         StringAttr::get(context, metadata.input_devices[i]));
    }
    for (int i : llvm::seq<int>(0, metadata.result_devices.size())) {
      func_op.setResultAttr(
          i, kTFDeviceAttr,
          StringAttr::get(context, metadata.result_devices[i]));
    }