| 26 | |
| 27 | template <typename Device, typename T, int NDIMS> |
| 28 | struct Slice { |
| 29 | void operator()(const Device& d, typename TTypes<T, NDIMS>::Tensor output, |
| 30 | typename TTypes<T, NDIMS>::ConstTensor input, |
| 31 | const Eigen::DSizes<Eigen::DenseIndex, NDIMS>& slice_indices, |
| 32 | const Eigen::DSizes<Eigen::DenseIndex, NDIMS>& slice_sizes) { |
| 33 | bool use_64bit = (input.size() > Eigen::NumTraits<int>::highest()); |
| 34 | if (!use_64bit && |
| 35 | Eigen::internal::is_same<Device, Eigen::GpuDevice>::value) { |
| 36 | Eigen::DSizes<int, NDIMS> indices; |
| 37 | for (int i = 0; i < NDIMS; ++i) { |
| 38 | indices[i] = slice_indices[i]; |
| 39 | } |
| 40 | Eigen::DSizes<int, NDIMS> sizes; |
| 41 | for (int i = 0; i < NDIMS; ++i) { |
| 42 | sizes[i] = slice_sizes[i]; |
| 43 | } |
| 44 | To32Bit(output).device(d) = To32Bit(input).slice(indices, sizes); |
| 45 | } else { |
| 46 | output.device(d) = input.slice(slice_indices, slice_sizes); |
| 47 | } |
| 48 | } |
| 49 | }; |
| 50 | |
| 51 | } // namespace functor |
| 52 | } // namespace tensorflow |
no outgoing calls