1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
|
/* Copyright 2015 Google Inc. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_KERNELS_WHERE_OP_H_
#define TENSORFLOW_KERNELS_WHERE_OP_H_
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#include "tensorflow/core/framework/tensor_types.h"
#include "tensorflow/core/platform/port.h"
namespace tensorflow {
namespace functor {
template <typename Device>
struct NumTrue {
EIGEN_ALWAYS_INLINE static void Compute(
const Device& d, typename TTypes<bool>::ConstFlat input,
TTypes<int64>::Scalar num_true) {
num_true.device(d) = input.template cast<int64>().sum();
}
};
template <typename Device, int NDIM>
struct Where {
EIGEN_ALWAYS_INLINE static void Compute(
const Device& d, typename TTypes<bool, NDIM>::ConstTensor input,
typename TTypes<int64>::Matrix output) {
Eigen::DenseIndex true_n = 0;
Eigen::DSizes<Eigen::DenseIndex, NDIM> dims = input.dimensions();
Eigen::DSizes<Eigen::DenseIndex, NDIM> strides;
// Calculate strides for RowMajor order.
EIGEN_STATIC_ASSERT((static_cast<int>(decltype(input)::Layout) ==
static_cast<int>(Eigen::RowMajor)),
INTERNAL_ERROR_INPUT_SHOULD_BE_ROWMAJOR);
strides[NDIM - 1] = 1;
for (int i = NDIM - 2; i >= 0; --i) {
strides[i] = strides[i + 1] * dims[i + 1];
}
// Note, no bounds checking is done on true_n. It is assumed that
// the output was correctly sized via output of NumTrue::Compute.
for (Eigen::DenseIndex n = 0; n < input.size(); ++n) {
if (input.data()[n]) {
WriteIndexRowMajor(output, strides, true_n, n);
++true_n;
}
}
}
EIGEN_ALWAYS_INLINE static void WriteIndexRowMajor(
typename TTypes<int64>::Matrix output,
const Eigen::DSizes<Eigen::DenseIndex, NDIM>& strides,
Eigen::DenseIndex true_n, Eigen::DenseIndex index) {
for (int i = 0; i < NDIM; ++i) {
output(true_n, i) = index / strides[i];
index %= strides[i];
}
}
};
} // namespace functor
} // namespace tensorflow
#endif // TENSORFLOW_KERNELS_WHERE_OP_H_
|
