Changing the type of a Eigen::Tensor(Map)

Question

I am currently building a custom TensorFlow Op. It is supposed to work like the the Conv2D-Op except it does that with a custom data type. As implementing a custom data type is relatively easy in Eigen and pretty hard in TensorFlow, I decided to copy the Eigen tensor to a new Eigen tensor with my custom datatype before TensorFlow calls Eigen. Converting the Eigen::TensorMap<Eigen::Tensor<float, 4, Eigen::RowMajor, Eigen::DenseIndex>, Eigen::Aligned> to a Eigen::TensorMap<Eigen::Tensor<CustomType, 4, Eigen::RowMajor, Eigen::DenseIndex>, Eigen::Aligned>, run the computation and then converting back to the floats.

I added some code in TensorFlows conv_2d.h, in the operator() of SpatialConvolution. I wrote two helper functions convertToCustomType and convertFromCustomType that are supposed to do the conversion for me. For the moment I don't really care about performance.

So basically I inject my two conversion functions before and after this line: https://github.com/tensorflow/tensorflow/blob/r1.12/tensorflow/core/kernels/conv_2d.h#L72

template<typename T>
Eigen::Tensor<CustomType, 4, Eigen::RowMajor, Eigen::DenseIndex> convertToCustomType(T &input) {
    Eigen::Tensor<CustomType, 4, Eigen::RowMajor, Eigen::DenseIndex> ret;
    ret.resize(input.dimensions());

    for (int a = 0; a < ret.size(); a++) {
        ret(a) = input(a);
    }
    return ret;
}

template<typename T1, typename T2>
void convertFromCustomType(T1 &input, T2 &output) {
    for (int a = 0; a < input.size(); a++) {
        output(a) = input(a);
    }
}

template<typename Device, typename T>
struct SpatialConvolution {
    void operator()(const Device &d, typename TTypes<T, 4>::Tensor output,
                    typename TTypes<T, 4>::ConstTensor input,
                    typename TTypes<T, 4>::ConstTensor filter, int row_stride,
                    int col_stride, int row_dilation, int col_dilation,
                    const Eigen::PaddingType &padding) {


        auto input_c = convertToCustomType(input);
        auto filter_c = convertToCustomType(filter);
        auto output_c = convertToCustomType(output);

        SpatialConvolutionFunc(d, output_c, input_c, filter_c, row_stride, col_stride, row_dilation, col_dilation, padding);

        convertFromCustomType(output_approx, output);
        output.device(d) = output;
    }
};

I also tried to run over all 4 dimensions of the tensor separately, that also doesn't appear to work.

template <typename T>
Eigen::Tensor<ApproxType, 4, Eigen::RowMajor> convertToCustomType(T input) {
    Eigen::Tensor<ApproxType, 4, Eigen::RowMajor> ret;
    ret.resize(input.dimensions());
    for (int a = 0; a < ret.dimension(0); a++) {
        for (int b = 0; b < ret.dimension(1); b++) {
            for (int c = 0; c < ret.dimension(2); c++) {
                for (int d = 0; d < ret.dimension(3); d++) {
                    ret(a, b, c, d) = input(a, b, c, d);
                }
            }
        }
    }
    return ret;
}

Both versions compile, but produce incorrect results. If I run my entire TensorFlow network with this custom Op it becomes non-deterministic, the output changes in distinct runs with the same input.

0
[[-0.06590138]]
1
[[-0.04544453]]
2
[[-0.0286443]]
3
[[-0.06590138]]
4
[[-0.06590138]]
5
[[-0.04544453]]

How am I supposed to change the actual type of an Eigen Tensor? I noticed there is something elegant like Tensor::cast<T>(), but calling that with T being anything else than Eigen::half does not compile. Am I maybe missing something in my custom type?

I know this is a pretty specific problem, but I would appreciate any ideas.

Answer 1

Apparently it is not enough to create the Tensor, it has to be initialized for example with ret.setZero() before filling it.