Searched refs:indices_nd (Results 1 – 4 of 4) sorted by relevance
64 const int indices_nd = SizeOfDimension(indices, indices_rank - 1); in Prepare() local73 if (indices_nd > params_rank) { in Prepare()84 const int output_rank = indices_rank + params_rank - indices_nd - 1; in Prepare()90 for (int i = indices_nd; i < params_rank; ++i) { in Prepare()
97 const int64 indices_nd = indices_shape.dim_size(indices_shape.dims() - 1); in DoGatherNd() local130 for (Index i = indices_nd; i < total_nd; ++i) { in DoGatherNd()165 switch (indices_nd) { in DoGatherNd()186 indices_nd); in DoGatherNd()196 gtl::ArraySlice<Index>(&indices_mat(bad_i, 0), indices_nd), ", "), in DoGatherNd()
1432 const int indices_nd = in ProcessGatherNdOperator() local1434 QCHECK_LE(indices_nd, input_shape.dimensions_count()); in ProcessGatherNdOperator()1440 for (int dim = indices_nd; dim < input_shape.dimensions_count(); ++dim) { in ProcessGatherNdOperator()
2983 const int indices_nd = indices_shape.Dims(indices_dims - 1); in GatherNd() local2988 for (int i = indices_nd; i < params_dims; ++i) { in GatherNd()2993 std::vector<int> dims_to_count(indices_nd, 0); in GatherNd()2994 for (int i = 0; i < indices_nd; ++i) { in GatherNd()3001 for (int j = 0; j < indices_nd; ++j) { in GatherNd()3002 from_pos += indices_data[i * indices_nd + j] * dims_to_count[j]; in GatherNd()