for (auto operand_type_and_idx : llvm::enumerate(op.getOperandTypes())) {
    Type operand_type = operand_type_and_idx.value();
    int index = operand_type_and_idx.index();
    tensorflow::tpu::TPUCompileMetadataProto::Arg* arg = metadata->add_args();
    tensorflow::DataType dtype;
    tensorflow::Status status =
        tensorflow::ConvertToDataType(operand_type, &dtype);
    if (!status.ok())
      return op.emitOpError(

}

Status ConvertAttribute(mlir::Type type, bool remove_ref_type,
                        AttrValue* value) {
  DataType dtype;
  TF_RETURN_IF_ERROR(ConvertToDataType(type, &dtype));
  if (tensorflow::IsRefType(dtype)) dtype = tensorflow::RemoveRefType(dtype);
  value->set_type(dtype);
  return absl::OkStatus();
}

Status ConvertAttribute(const mlir::TypeAttr& type, bool remove_ref_type,

  %0 = "tf.Const"() {device = "/device:CPU:0", value = dense<9> : tensor<i32>} : () -> tensor<i32>
  %1 = "tf.ReadVariableOp"(%arg) {device = "/device:CPU:0", dtype = i32} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
  %2 = "tf.Less"(%1, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
  func.return %2 : tensor<i1>
}

    }

    se::Stream* const src_compute_stream = src_xla_context->stream();
    TF_RET_CHECK(src_compute_stream != nullptr);
    TF_RET_CHECK(input->dtype() == output->dtype())
        << "input type: " << DataTypeString(input->dtype()) << " output type "
        << DataTypeString(output->dtype());
    TF_RET_CHECK(input->shape() == output->shape());
    TF_RET_CHECK(DMAHelper::CanUseDMA(input));
    auto* const src_compute_stream_impl =

                 absl::Span<const int64_t> shape, const TF_DataType dtype)
      : device_(device),
        tensors_(std::move(tensors)),
        shape_(std::vector<int64_t>(shape.begin(), shape.end())),
        dtype_(dtype) {}
  ParallelTensor(const ParallelDevice& device,
                 std::vector<TensorHandlePtr> tensors, const TF_DataType dtype)
      : device_(device),
        tensors_(std::move(tensors)),

      // Return values
      {"y: float"},
      // Attr def
      {},
      // Nodes
      {
          {{"two"},
           "Const",
           {},
           {{"value", test::AsScalar<int32>(2)}, {"dtype", DT_INT64}}},
          {{"scale"},
           "Cast",
           {"two"},
           {{"SrcT", DT_INT64}, {"DstT", DT_FLOAT}}},
          {{"y"}, "Mul", {"x", "scale"}, {{"T", DT_FLOAT}}},
      });

                                     const char* attr_name,
                                     const TF_DataType type,
                                     TF_Status* status) {
  tensorflow::DataType dtype = static_cast<tensorflow::DataType>(type);
  tensorflow::AddTypeConstraint(kernel_builder, attr_name, dtype, status);
}

void TF_KernelBuilder_HostMemory(TF_KernelBuilder* kernel_builder,

      const Tensor* value = variable.var()->tensor();
      arg.type = value->dtype();
      arg.shape = value->shape();
      arg.initialized = true;
    } else {
      // Instead of embedding constant into HLO,
      // we handle tf.constant as parameter to reduce size.
      arg.kind = XlaCompiler::Argument::kParameter;
      arg.type = input->dtype();
      arg.shape = input->shape();
    }
    args.push_back(arg);

  }

  tensorflow::DataType dtype;
  auto status = tensorflow::ConvertToDataType(operand.getType(), &dtype);
  if (!status.ok()) {
    op->emitRemark() << "skipping legalization due to " << status.ToString();
    return tensorflow::XlaExpression::Invalid();
  }
  return tensorflow::XlaExpression::XlaOp(xla_op, dtype);
}

}  // namespace mhlo

  ctx.appendDialectRegistry(registry);
  ctx.loadAllAvailableDialects();
}

Status ConvertDataTypeToTensor(tensorflow::DataType dtype, Builder builder,
                               Type* type) {
  Status s = tensorflow::ConvertDataType(dtype, builder, type);
  if (s.ok()) *type = UnrankedTensorType::get(*type);
  return s;
}

class MlirTensor : public TracingTensorHandle {
 public:
  explicit MlirTensor(Value value)