for (int i = 0; i < 100; ++i) {
    TensorHandlePtr multiply_result(
        Multiply(context.get(), combined_value.get(), combined_value.get(),
                 status.get()));
    ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
    std::array<TensorHandlePtr, 2> out_components;
    ExtractPerDeviceValues(context.get(), multiply_result.get(),
                           &out_components, status.get());

  } else {
    TF_SetStatus(
        status, TF_UNIMPLEMENTED,
        absl::StrCat(
            "Trying to copy a tensor out of a parallel device. Since there "
            "are multiple components to parallel tensors, they must be "
            "unpacked explicitly.\n",
            tensorflow::unwrap(tensor)->DebugString())
            .c_str());
    return nullptr;
  }
}

// Allows a single op at a time to be launched without blocking.
//
// DeviceThread itself is thread-safe, in that StartExecute will block if there
// is a pending execution. Since StartExecute is equivalent to grabbing a lock,
// multiple DeviceThreads should always be accessed in the same order to avoid
// deadlocks.
class DeviceThread {
 public:
  // Starts a background thread waiting for `StartExecute`.

      context.get(), components, second_device_name, status.get());
  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());

  TensorHandlePtr multiply_result(
      Multiply(context.get(), second_combined_value.get(),
               second_negative_one.get(), status.get()));
  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());

  if (TF_GetCode(status) != TF_OK) return;
  for (int i = 0; i < num_replicas; ++i) {
    (*components)[i].reset(result_handles[i]);
  }
}

TensorHandlePtr Multiply(TFE_Context* context, TFE_TensorHandle* first,
                         TFE_TensorHandle* second, TF_Status* status) {
  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> op(
      TFE_NewOp(context, "Mul", status), TFE_DeleteOp);

// How to upload new data when Flush() is called multiple times.
// By default the entire file is reuploaded.
constexpr char kAppendMode[] = "GCS_APPEND_MODE";
// If GCS_APPEND_MODE=compose then instead the new data is uploaded to a
// temporary object and composed with the original object. This is disabled by
// default as the multiple API calls required add a risk of stranding temporary
// objects.

//  --test_arg=--scheme= --test_arg=--scheme=file
//
// Note that to test the local filesystem we use an empty value.

namespace tensorflow {
namespace {

using ::tensorflow::error::Code;

// As we need to test multiple URI schemes we need a parameterized test.
// Furthermore, since each test creates and deletes files, we will use the same
// fixture to create new directories in `SetUp`. Each directory will reside in

TEST_P(UnifiedCAPI, TestBasicGraphMatMul) {
  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
      TF_NewStatus(), TF_DeleteStatus);

  // Start a new function / execution context.
  string fn_name = "matrix_multiply";
  TF_ExecutionContext* graph_ctx =
      TF_CreateFunction(fn_name.c_str(), status.get());
  ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());

  auto* placeholder_t =