std::optional<xla::DeviceAssignmentProto> xla_device_assignment;
};

// Extracts device coordinates from a device assignment attribute on an op.
absl::StatusOr<llvm::SmallVector<int64_t, 8>> GetDeviceCoordinates(
    mlir::ArrayAttr device_assignment_attr);

// Finds the TPU compilation device and execution devices from `devices` for a
// TPU computation subgraph. Compilation device is determined from looking up

                                   .str());
  if (absl::StatusOr<llvm::SmallVector<int64_t, 8>> fetched_device_coordinates =
          tensorflow::GetDeviceCoordinates(device_assignment_attr);
      fetched_device_coordinates.ok()) {
    device_coordinates = *fetched_device_coordinates;
    return mlir::success();
  } else {

                                             tensorflow::kDeviceAssignmentAttr)
                                   .str());

  auto status_or_device_coodinates =
      tensorflow::GetDeviceCoordinates(device_assignment_attr);
  if (!status_or_device_coodinates.ok())
    return cluster.emitError()
           << "error in fetching tpu device coordinates: "
           << status_or_device_coodinates.status().message();

  mlir::MLIRContext context;
  mlir::Builder builder(&context);
  auto device_assignment_attr = builder.getI64ArrayAttr({1, 2, 3});
  auto status_or_device_coodinates =
      GetDeviceCoordinates(device_assignment_attr);
  ASSERT_TRUE(status_or_device_coodinates.ok());
  auto device_coordinates = status_or_device_coodinates.value();
  EXPECT_EQ(device_coordinates[0], 1);
  EXPECT_EQ(device_coordinates[1], 2);

        llvm::formatv("requires attribute '{0}'",
                      tensorflow::kDeviceAssignmentAttr)
            .str());

  auto status_or_device_coodinates =
      tensorflow::GetDeviceCoordinates(device_assignment_attr);
  if (!status_or_device_coodinates.ok())
    return cluster_func.emitError()
           << "error in fetching tpu device coordinates: "