diff options
| author |  Benoit Steiner <benoit.steiner.goog@gmail.com> | 2015-08-31 08:18:53 -0700 |
|---|---|---|
| committer | Benoit Steiner <benoit.steiner.goog@gmail.com> | 2015-08-31 08:18:53 -0700 |
| commit | f41831e445f3fdd9dc324561135b2a19eafd9a56 (patch) | |
| tree | 045cb917d62685b342ce129e384e03e63c916898 /unsupported/Eigen/CXX11/src/Tensor/TensorBase.h | |
| parent | 2ab603316af7c1bcf1d5e87d9ba50a2589b36e37 (diff) | |
Added support for argmax/argmin
Diffstat (limited to 'unsupported/Eigen/CXX11/src/Tensor/TensorBase.h')
| -rw-r--r-- | unsupported/Eigen/CXX11/src/Tensor/TensorBase.h | 59 |
1 files changed, 59 insertions, 0 deletions
diff --git a/unsupported/Eigen/CXX11/src/Tensor/TensorBase.h b/unsupported/Eigen/CXX11/src/Tensor/TensorBase.h index 0e5e4b426..477e4a174 100644 --- a/unsupported/Eigen/CXX11/src/Tensor/TensorBase.h +++ b/unsupported/Eigen/CXX11/src/Tensor/TensorBase.h @@ -363,6 +363,58 @@ class TensorBase<Derived, ReadOnlyAccessors> return TensorReductionOp<internal::MinReducer<CoeffReturnType>, const DimensionList<Index, NumDimensions>, const Derived>(derived(), in_dims, internal::MinReducer<CoeffReturnType>()); } + EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE + const TensorTupleReducerOp< + internal::ArgMaxTupleReducer<Tuple<Index, CoeffReturnType> >, + const array<Index, NumDimensions>, const Derived> + argmax() const { + array<Index, NumDimensions> in_dims; + for (int d = 0; d < NumDimensions; ++d) in_dims[d] = d; + return TensorTupleReducerOp< + internal::ArgMaxTupleReducer<Tuple<Index, CoeffReturnType> >, + const array<Index, NumDimensions>, + const Derived>(derived(), internal::ArgMaxTupleReducer<Tuple<Index, CoeffReturnType> >(), -1, in_dims); + } + + EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE + const TensorTupleReducerOp< + internal::ArgMinTupleReducer<Tuple<Index, CoeffReturnType> >, + const array<Index, NumDimensions>, const Derived> + argmin() const { + array<Index, NumDimensions> in_dims; + for (int d = 0; d < NumDimensions; ++d) in_dims[d] = d; + return TensorTupleReducerOp< + internal::ArgMinTupleReducer<Tuple<Index, CoeffReturnType> >, + const array<Index, NumDimensions>, + const Derived>(derived(), internal::ArgMinTupleReducer<Tuple<Index, CoeffReturnType> >(), -1, in_dims); + } + + EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE + const TensorTupleReducerOp< + internal::ArgMaxTupleReducer<Tuple<Index, CoeffReturnType> >, + const array<Index, 1>, const Derived> + argmax(const int return_dim) const { + array<Index, 1> in_dims; + in_dims[0] = return_dim; + return TensorTupleReducerOp< + internal::ArgMaxTupleReducer<Tuple<Index, CoeffReturnType> >, + const array<Index, 1>, + const Derived>(derived(), internal::ArgMaxTupleReducer<Tuple<Index, CoeffReturnType> >(), return_dim, in_dims); + } + + EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE + const TensorTupleReducerOp< + internal::ArgMinTupleReducer<Tuple<Index, CoeffReturnType> >, + const array<Index, 1>, const Derived> + argmin(const int return_dim) const { + array<Index, 1> in_dims; + in_dims[0] = return_dim; + return TensorTupleReducerOp< + internal::ArgMinTupleReducer<Tuple<Index, CoeffReturnType> >, + const array<Index, 1>, + const Derived>(derived(), internal::ArgMinTupleReducer<Tuple<Index, CoeffReturnType> >(), return_dim, in_dims); + } + template <typename Reducer, typename Dims> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const TensorReductionOp<Reducer, const Dims, const Derived> reduce(const Dims& dims, const Reducer& reducer) const { @@ -483,6 +535,13 @@ class TensorBase<Derived, ReadOnlyAccessors> return TensorInflationOp<const Strides, const Derived>(derived(), strides); } + // Returns a tensor containing index/value tuples + EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE + const TensorIndexTupleOp<const Derived> + index_tuples() const { + return TensorIndexTupleOp<const Derived>(derived()); + } + // Support for custom unary and binary operations template <typename CustomUnaryFunc> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE |