| 27 | // Functor used by PadOp to do the computations. |
| 28 | template <typename Device, typename T, typename Tpadding, int Dims> |
| 29 | struct Pad { |
| 30 | // Pad "input" into "output", as specified by "paddings" and "pad_value". |
| 31 | // See pad_op.cc for details. |
| 32 | void operator()(const Device& d, typename TTypes<T, Dims>::Tensor output, |
| 33 | typename TTypes<T, Dims>::ConstTensor input, |
| 34 | Eigen::array<Eigen::IndexPair<Tpadding>, Dims> paddings, |
| 35 | T pad_value) { |
| 36 | if (Eigen::internal::is_same<Device, Eigen::GpuDevice>::value && |
| 37 | (output.size() <= std::numeric_limits<int32>::max())) { |
| 38 | To32Bit(output).device(d) = To32Bit(input).pad(paddings, pad_value); |
| 39 | } else { |
| 40 | output.device(d) = input.pad(paddings, pad_value); |
| 41 | } |
| 42 | } |
| 43 | }; |
| 44 | |
| 45 | template <typename Device, typename T, typename Tpadding> |
| 46 | struct Pad<Device, T, Tpadding, 0> { |
no outgoing calls