Should copyFunc in TF_AssignVariable do a device memcpy of the buffer?


on my way to reconstructing Tensorflow’s internal puzzle for plugin development, I am having a doubt on what the copyFunc parameter of TF_AssignVariable should do.

Should it perform a copy of the data from and to the underlying buffers?
Should it call another tensorflow mechanism like forward input, bitcast, etc?

As a reminder, here is the signature:

// Expose higher level Assignment operation for Pluggable vendors to implementkernels_experi
// in the plugin for Training. The API takes in the context with indices for
// the input and value tensors. It also accepts the copy callback provided by
// pluggable vendor to do the copying of the tensors. The caller takes ownership
// of the `source` and `dest` tensors and is responsible for freeing them with
// TF_DeleteTensor. This function will return an error when the following
// conditions are met:
//   1. `validate_shape` is set to `true`
//   2. The variable is initialized
//   3. The shape of the value tensor doesn't match the shape of the variable
//      tensor.
TF_CAPI_EXPORT extern void TF_AssignVariable(
    TF_OpKernelContext* ctx, int input_index, int value_index,
    bool validate_shape,
    void (*copyFunc)(TF_OpKernelContext* ctx, TF_Tensor* source,
                     TF_Tensor* dest),
    TF_Status* status);