// 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,
    std::array<TensorHandlePtr, num_replicas>* components, TF_Status* status);

// Helper to pack `num_replicas` TFE_TensorHandles into one parallel handle.
template <std::size_t num_replicas>
TensorHandlePtr CreatePerDeviceValues(

  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,
    std::array<TensorHandlePtr, num_replicas>* components, TF_Status* status) {
  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> op(

  TFE_TensorHandle* packed_handle = nullptr;
  if (has_packed_input) {
    int num_replicas = 1;
    std::vector<TFE_TensorHandle*> packed_handles = {h1_task2};
    packed_handle = TFE_CreatePackedTensorHandle(ctx, packed_handles.data(),
                                                 &num_replicas, status);
    ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
  }

  TFE_Op* matmul = nullptr;

// parallel device creates a parallel tensor `x` with `a` on the first of
// `underlying_devices` and `b` on the second. Running `a_unpacked, b_unpacked =
// TPUReplicatedOutput(input=x, num_replicas=2)` un-packs the parallel tensor
// into its components.
//
// The filled `device` struct and the allocated `device_info` struct may be
// passed to TFE_RegisterCustomDevice. The `device_name` arguments must match.

  // placed on task1, Handle types are LOCAL, REMOTE, LOCAL on task1.
  int num_replicas = 3;
  std::vector<TFE_TensorHandle*> handles = {h0, h1, h2};
  TFE_TensorHandle* packed_handle =
      TFE_CreatePackedTensorHandle(ctx, handles.data(), &num_replicas, status);
  ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
  EXPECT_EQ(TFE_TensorHandleDataType(packed_handle), TF_RESOURCE);

        TFE_NewOp(context.get(), "TPUReplicatedOutput", status.get()),
        TFE_DeleteOp);
    if (TF_GetCode(status.get()) != TF_OK) return;
    TFE_OpSetAttrInt(op.get(), "num_replicas", 1);
    TFE_OpAddInput(op.get(), value_one.get(), status.get());
    if (TF_GetCode(status.get()) != TF_OK) return;
    TFE_OpSetDevice(op.get(), device_name, status.get());