diff options
| author |  Antonio Sanchez <cantonios@google.com> | 2021-06-11 08:30:41 -0700 |
|---|---|---|
| committer |  Rasmus Munk Larsen <rmlarsen@google.com> | 2021-06-11 17:02:19 +0000 |
| commit | 6aec83263d32c29f6c5623b9716ec7e367693078 (patch) | |
| tree | 2b2129cb764bfaae50bce291e9b573cc82c94884 | |
| parent | fc87e2cbaa65e7e93a2c695ce5a9dc048a64a985 (diff) | |
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.
| -rw-r--r-- | unsupported/Eigen/CXX11/src/Tensor/TensorContraction.h | 4 |
1 files changed, 4 insertions, 0 deletions
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<const Derived*>(this)->template evalProduct<Unaligned>(buffer); |