blob: c4f8a4abdab328bed93d3db9892cce847b504345 (
plain)
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
|
#ifndef TENSORFLOW_KERNELS_PAD_OP_H_
#define TENSORFLOW_KERNELS_PAD_OP_H_
// Functor definition for PadOp, must be compilable by nvcc.
#include "tensorflow/core/platform/port.h"
#include "tensorflow/core/framework/tensor_types.h"
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
namespace tensorflow {
namespace functor {
// Functor used by PadOp to do the computations.
template <typename Device, typename T, int Dims>
struct Pad {
// Pad "input" into "output", as specified by "paddings". See pad_op.cc for
// details.
void operator()(const Device& d, typename TTypes<T, Dims>::Tensor output,
typename TTypes<T, Dims>::ConstTensor input,
Eigen::array<std::pair<int32, int32>, Dims> paddings) {
output.device(d) = input.pad(paddings);
}
};
} // namespace functor
} // namespace tensorflow
#endif // TENSORFLOW_KERNELS_PAD_OP_H_
|
