| 47 | |
| 48 | template <typename Device, typename T, typename Index> |
| 49 | Status DoGatherNd(OpKernelContext* c, const Tensor& params, |
| 50 | const Tensor& indices, Tensor* out) { |
| 51 | if (!TensorShapeUtils::IsVectorOrHigher(params.shape())) { |
| 52 | return errors::InvalidArgument("params must be at least a vector"); |
| 53 | } |
| 54 | if (!TensorShapeUtils::IsVectorOrHigher(indices.shape())) { |
| 55 | return errors::InvalidArgument("indices must be at least a vector"); |
| 56 | } |
| 57 | if (indices.dim_size(indices.dims() - 1) > params.dims()) { |
| 58 | return errors::InvalidArgument( |
| 59 | "index innermost dimension length must be <= params rank; saw: ", |
| 60 | indices.dim_size(indices.dims() - 1), " vs. ", params.dims()); |
| 61 | } |
| 62 | |
| 63 | const TensorShape& indices_shape(indices.shape()); |
| 64 | const int64 indices_nd = indices_shape.dim_size(indices_shape.dims() - 1); |
| 65 | |
| 66 | // Check that we have enough index space |
| 67 | int64 N_big = 1; |
| 68 | for (int i = 0; i < indices_shape.dims() - 1; ++i) { |
| 69 | N_big *= indices_shape.dim_size(i); |
| 70 | } |
| 71 | if (N_big > std::numeric_limits<int>::max()) { |
| 72 | return errors::InvalidArgument( |
| 73 | "indices has too many elements for int indexing: ", N_big, " > ", |
| 74 | std::numeric_limits<int>::max()); |
| 75 | } |
| 76 | if (params.NumElements() > std::numeric_limits<Index>::max()) { |
| 77 | return errors::InvalidArgument("params.NumElements() too large for ", |
| 78 | DataTypeString(DataTypeToEnum<Index>::v()), |
| 79 | " indexing: ", params.NumElements(), " > ", |
| 80 | std::numeric_limits<Index>::max()); |
| 81 | } |
| 82 | |
| 83 | // The result shape is |
| 84 | // indices.shape[:-1] + params.shape[indices.shape[-1]:] |
| 85 | Index N_result = 1; |
| 86 | for (int i = 0; i < indices_shape.dims() - 1; ++i) { |
| 87 | N_result *= indices_shape.dim_size(i); |
| 88 | } |
| 89 | |
| 90 | const TensorShape& params_shape(params.shape()); |
| 91 | Index total_nd = params_shape.dims(); |
| 92 | |
| 93 | TensorShape result_shape(indices_shape); |
| 94 | result_shape.RemoveLastDims(1); |
| 95 | |
| 96 | int64 slice_size_big = 1; |
| 97 | for (Index i = indices_nd; i < total_nd; ++i) { |
| 98 | slice_size_big *= params_shape.dim_size(i); |
| 99 | result_shape.AddDim(params_shape.dim_size(i)); |
| 100 | } |
| 101 | |
| 102 | if (slice_size_big > std::numeric_limits<Index>::max()) { |
| 103 | return errors::InvalidArgument( |
| 104 | "slice size is too large for indexing: ", slice_size_big, " > ", |
| 105 | std::numeric_limits<Index>::max()); |
| 106 | } |
nothing calls this directly
no test coverage detected