TensorFlowカスタムオペレーションで2つの行列を乗算するにはどうすればよいですか？

Question

私はtensorflow/core/kernelsの例からまとめた次のコードで2つの行列を掛けようとしています。

A = [[ 25.90848866 23.31998387 38.59424611] 
    [ 48.12131663 48.15446712 0.27323556] 
    [ 50.99763646 25.77013947 24.67779583] 
    [ 14.43187796 20.18683861 45.63747844] 
    [ 23.09759944 49.91747332 31.83040427] 
    [ 3.01687455 15.14876558 41.13147952] 
    [ 41.32467959 49.91440187 21.21775212] 
    [ 5.48660952 39.98743312 3.33354911]] 

B = [[ 0.999994 1.98149e-05 -0.00337665 ] 
    [ 1.98149e-05 0.999931 0.011736 ] 
    [ 0.00337665 -0.011736 0.999925 ]] 

で

void Compute(OpKernelContext* context) override { 

    const Tensor& A = context->input(0); 
    const Tensor& B = context->input(1); 

    const auto A_mat = A.matrix<T>(); 
    const auto B_mat = B.matrix<T>(); 

    Tensor* output_tensor = nullptr; 
    OP_REQUIRES_OK(context, context->allocate_output(0, A_tensor.shape(), &output_tensor)); 
    auto output = output_tensor->matrix<T>(); 

    output = A_mat * B_mat; 
} 

opがoutput = A_mat + A_mat（プラス代わりMATMULの）である場合、私は、しかし

output = [[ 4.04171173e+001 6.19674481e-004 -6.08213832e-002] 
      [ 6.27670070e-004 5.15969690e+000 4.23702253e-001] 
      [ 1.18978160e-001 -3.34875362e-001 1.78214391e+001] 
      [ 0.00000000e+000 1.78005398e-313 2.44005345e-309] 
      [ 4.62065626e+064 0.00000000e+000 0.00000000e+000] 
      [    nan    nan 5.00303243e+065] 
      [ 3.45845952e-323 2.02072849e-321 2.72281224e-308] 
      [ 0.00000000e+000 3.28336636e-308 0.00000000e+000]] 

を取得し、出力は予想通り2*A_matです。ここで何が起こっているのですか？カスタム演算でこのような固有ベクトル、行列、テンソルを扱うことはできませんか？私は何が欠けていますか？

Answer 1

how to change 2D Eigen::Tensor to Eigen::Matrixのおかげで、2つのテンソルにEigen表記を乗算することは簡単でした。

const auto A_mat = A.matrix<T>();に関連する問題はtensorflow::Tensorを返し、Eigen::Matrixを返しません。これは、以下に示すようにEigen::Mapで実現できます。

void Compute(OpKernelContext* context) override { 

    const Tensor& A = context->input(0); 
    const Tensor& B = context->input(1); 

    const auto A_mat = Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(
     A.flat<T>().data(), A.dim_size(0), A.dim_size(1)); 

    const auto B_mat = Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(
     B.flat<T>().data(), B.dim_size(0), B.dim_size(1)); 

    Tensor* output_tensor = nullptr; 
    OP_REQUIRES_OK(context, context->allocate_output(0, A_tensor.shape(), &output_tensor)); 

    const auto output = Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(
     output.flat<T>().data(), output.dim_size(0), output.dim_size(1)); 

    output = A_mat * B_mat; 
} 

私はCUDAでセグメンテーション違反を取得し、この後