aboutsummaryrefslogtreecommitdiffhomepage
path: root/unsupported/Eigen/CXX11/src/Tensor/TensorBase.h
diff options
context:
space:
mode:
authorGravatar Eugene Zhulenev <ezhulenev@google.com>2018-07-10 13:16:38 -0700
committerGravatar Eugene Zhulenev <ezhulenev@google.com>2018-07-10 13:16:38 -0700
commit01fd4096d395e7b816459f571bf2328c8435cc37 (patch)
tree02b928b34f77c3e63126c3175b6ea06174818f51 /unsupported/Eigen/CXX11/src/Tensor/TensorBase.h
parent5539587b1f5b5922b2419b0a4468cf2f393def51 (diff)
Fuse computations into the Tensor contractions using output kernel
Diffstat (limited to 'unsupported/Eigen/CXX11/src/Tensor/TensorBase.h')
-rw-r--r--unsupported/Eigen/CXX11/src/Tensor/TensorBase.h10
1 files changed, 8 insertions, 2 deletions
diff --git a/unsupported/Eigen/CXX11/src/Tensor/TensorBase.h b/unsupported/Eigen/CXX11/src/Tensor/TensorBase.h
index bdc1a17a7..97f90f638 100644
--- a/unsupported/Eigen/CXX11/src/Tensor/TensorBase.h
+++ b/unsupported/Eigen/CXX11/src/Tensor/TensorBase.h
@@ -517,9 +517,15 @@ class TensorBase<Derived, ReadOnlyAccessors>
typedef Eigen::IndexPair<Index> DimensionPair;
template<typename OtherDerived, typename Dimensions> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
- const TensorContractionOp<const Dimensions, const Derived, const OtherDerived>
+ const TensorContractionOp<const Dimensions, const Derived, const OtherDerived, const NoOpOutputKernel>
contract(const OtherDerived& other, const Dimensions& dims) const {
- return TensorContractionOp<const Dimensions, const Derived, const OtherDerived>(derived(), other.derived(), dims);
+ return TensorContractionOp<const Dimensions, const Derived, const OtherDerived, const NoOpOutputKernel>(derived(), other.derived(), dims);
+ }
+
+ template<typename OtherDerived, typename Dimensions, typename OutputKernel> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
+ const TensorContractionOp<const Dimensions, const Derived, const OtherDerived, const OutputKernel>
+ contract(const OtherDerived& other, const Dimensions& dims, const OutputKernel& output_kernel) const {
+ return TensorContractionOp<const Dimensions, const Derived, const OtherDerived, const OutputKernel>(derived(), other.derived(), dims, output_kernel);
}
// Convolutions.