特征::张量或张量流::张量应在张量流自定义操作中使用？

Question

我在使用cuda的张量流中添加了一个自定义操作，大多数计算是许多小矩阵运算（matmul，inverse，...），我使用浮点型数据和cublas函数通过cuda内核函数实现了此操作。似乎使用tensorflow :: tensor操作可能是一个更好的选择。我的问题是在操作tensorflow :: Tensor作为输入时，应该将其转换为eigen :: tensor并使用本征函数，或者如果tf具有此c ++接口，则直接使用tensorflow :: tensor函数（例如矩阵mul）？

我实现的代码如下

void Compute(OpKernelContext* context) override
{
    const Tensor& a_tensor = context->input(0);
    auto a = a_tensor.flat<complex64>();
    kenerlLauncher(a, ...);
}

void kenerlLauncher(const TTypes<complex64>::ConstFlat a)
{
    g_func << <grid, threads >> > ((float*)a.data(), ...);
}

__global__ void g_func(const float* a, ...)
{
    // ...
}

如果使用张量，则代码可能为

void Compute(OpKernelContext* context) override
{
    const Tensor& a_tensor = context->input(0);
    auto a = a_tensor.flat<complex64, 4>();
    kenerlLauncher(a, ...);
}

void kenerlLauncher(const TTypes<complex64, 4>::ConstTensor a)
{
    g_func << <grid, threads >> > (a, ...);
}

__global__ void g_func(const TTypes<complex64, 4>::ConstTensor a, ...)
{
    // a has to be converted to eigen::tensor ?
    // eigen::tensor can be used in cuda kernel function ?
    // or tensorflow c++ interface of matrix op should be used ?
}