// A ParallelDevice on its own is not registered with a TFE_Context, and so has
// no device name (e.g. for `tf.device`). `NamedParallelDevice` associates a
// name with it, which lets us pack its `ParallelTensor`s into TFE_TensorHandles
// placed on the parallel device.
class NamedParallelDevice {
 public:
  NamedParallelDevice(const std::string& name,

      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());

  // Un-pack the parallel tensor to verify that the operation was
  // successful. The resulting structure should be:
  //   second_device{first_device{1. * 3., 2. * 3.}, 3. * 3.}.
  std::array<TensorHandlePtr, 2> second_components;

                   &FloatDeallocator, nullptr),
      TF_DeleteTensor);
  return TensorHandlePtr(TFE_NewTensorHandle(tensor.get(), status));
}

// Helper to un-pack `num_replicas` TFE_TensorHandles from one parallel handle.
template <std::size_t num_replicas>
void ExtractPerDeviceValues(
    TFE_Context* context, TFE_TensorHandle* input,

  if (first_bad_status != nullptr) {
    TF_SetStatus(status, TF_GetCode(first_bad_status.get()),
                 TF_Message(first_bad_status.get()));
    return result;
  }
  // For each output of the original operation, pack the per-device
  // TensorHandles we've computed into a single parallel TensorHandle.
  std::vector<std::unique_ptr<ParallelTensor>> per_device_outputs;
  per_device_outputs.reserve(first_op_output_count);

  // Add a sync point to make sure that variables have been initialized
  // before the function execution starts.
  TFE_ContextAsyncWait(ctx, status);
  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);

  // Pack 3 variable handles into one TFE_TensorHandle.
  // When remote is false, function device is placed on task0. Handle types are
  // REMOTE, REMOTE, LOCAL on task0. When remote is true, function device is

        TFE_DeleteTensorHandle(outputs[i]);
      }
    }
    return status.status;
  }

  tensorflow::Status Pack(absl::Span<ImmediateExecutionTensorHandle*> handles,
                          ImmediateExecutionTensorHandle** result) override {
    TF_Status status;
    *result = tensorflow::unwrap(device_.pack(context_,
                                              tensorflow::wrap(handles.data()),

      // One of the inputs we're trying to pack is on a custom device. We'll let
      // the first custom device we see handle all of the packing.
      auto* custom_device_handle =
          tensorflow::down_cast<tensorflow::CustomDeviceTensorHandle*>(
              unwrapped_handle);
      tensorflow::ImmediateExecutionTensorHandle* result;
      status->status = custom_device_handle->device()->Pack(