From 6aec83263d32c29f6c5623b9716ec7e367693078 Mon Sep 17 00:00:00 2001 From: Antonio Sanchez Date: Fri, 11 Jun 2021 08:30:41 -0700 Subject: [PATCH] Allow custom TENSOR_CONTRACTION_DISPATCH macro. Currently TF lite needs to hack around with the Tensor headers in order to customize the contraction dispatch method. Here we add simple `#ifndef` guards to allow them to provide their own dispatch prior to inclusion. --- unsupported/Eigen/CXX11/src/Tensor/TensorContraction.h | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/unsupported/Eigen/CXX11/src/Tensor/TensorContraction.h b/unsupported/Eigen/CXX11/src/Tensor/TensorContraction.h index d442c782c..8b35f7985 100644 --- a/unsupported/Eigen/CXX11/src/Tensor/TensorContraction.h +++ b/unsupported/Eigen/CXX11/src/Tensor/TensorContraction.h @@ -633,6 +633,7 @@ struct TensorContractionEvaluatorBase : internal::no_assignment_operator } #endif // EIGEN_USE_THREADS +#ifndef TENSOR_CONTRACTION_DISPATCH #define TENSOR_CONTRACTION_DISPATCH(METHOD, ALIGNMENT, ARGS) \ if (this->m_lhs_inner_dim_contiguous) { \ if (this->m_rhs_inner_dim_contiguous) { \ @@ -663,7 +664,9 @@ struct TensorContractionEvaluatorBase : internal::no_assignment_operator } \ } \ } +#endif +#ifndef TENSOR_CONTRACTION_ASYNC_DISPATCH #define TENSOR_CONTRACTION_ASYNC_DISPATCH(METHOD, DONE, ALIGNMENT, ARGS, FN) \ if (this->m_lhs_inner_dim_contiguous) { \ if (this->m_rhs_inner_dim_contiguous) { \ @@ -694,6 +697,7 @@ struct TensorContractionEvaluatorBase : internal::no_assignment_operator } \ } \ } +#endif EIGEN_DEVICE_FUNC void evalTo(Scalar* buffer) const { static_cast(this)->template evalProduct(buffer);