Hi,
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);