| 288 | } |
| 289 | |
| 290 | void Compute(OpKernelContext* context) override { |
| 291 | const Tensor& points_tensor = context->input(0); |
| 292 | const Tensor& centers_tensor = context->input(1); |
| 293 | const Tensor& k_tensor = context->input(2); |
| 294 | |
| 295 | OP_REQUIRES(context, TensorShapeUtils::IsMatrix(points_tensor.shape()), |
| 296 | InvalidArgument("Input points should be a matrix.")); |
| 297 | OP_REQUIRES(context, TensorShapeUtils::IsMatrix(centers_tensor.shape()), |
| 298 | InvalidArgument("Input centers should be a matrix.")); |
| 299 | OP_REQUIRES(context, TensorShapeUtils::IsScalar(k_tensor.shape()), |
| 300 | InvalidArgument("Input k should be a scalar.")); |
| 301 | |
| 302 | const int64 num_points = points_tensor.dim_size(0); |
| 303 | const int64 point_dimensions = points_tensor.dim_size(1); |
| 304 | const int64 num_centers = centers_tensor.dim_size(0); |
| 305 | const int64 center_dimensions = centers_tensor.dim_size(1); |
| 306 | |
| 307 | OP_REQUIRES(context, num_points > 0, |
| 308 | InvalidArgument("Expected points.rows() > 0.")); |
| 309 | OP_REQUIRES( |
| 310 | context, point_dimensions == center_dimensions, |
| 311 | InvalidArgument("Expected point_dimensions == center_dimensions: ", |
| 312 | point_dimensions, " vs ", center_dimensions, ".")); |
| 313 | |
| 314 | const Eigen::Map<const MatrixXfRowMajor> points( |
| 315 | points_tensor.matrix<float>().data(), num_points, point_dimensions); |
| 316 | const Eigen::Map<const MatrixXfRowMajor> centers( |
| 317 | centers_tensor.matrix<float>().data(), num_centers, center_dimensions); |
| 318 | const int64 k = std::min<int64>(num_centers, k_tensor.scalar<int64>()()); |
| 319 | |
| 320 | Tensor* output_nearest_center_indices_tensor; |
| 321 | Tensor* output_nearest_center_distances_tensor; |
| 322 | OP_REQUIRES_OK(context, context->allocate_output( |
| 323 | 0, TensorShape({num_points, k}), |
| 324 | &output_nearest_center_indices_tensor)); |
| 325 | OP_REQUIRES_OK(context, context->allocate_output( |
| 326 | 1, TensorShape({num_points, k}), |
| 327 | &output_nearest_center_distances_tensor)); |
| 328 | |
| 329 | if (k == 0) return; |
| 330 | |
| 331 | Eigen::Map<MatrixXi64RowMajor> nearest_center_indices( |
| 332 | output_nearest_center_indices_tensor->matrix<int64>().data(), |
| 333 | num_points, k); |
| 334 | Eigen::Map<MatrixXfRowMajor> nearest_center_distances( |
| 335 | output_nearest_center_distances_tensor->matrix<float>().data(), |
| 336 | num_points, k); |
| 337 | |
| 338 | const Eigen::VectorXf centers_half_squared_norm = |
| 339 | 0.5 * centers.rowwise().squaredNorm(); |
| 340 | |
| 341 | // The distance computation is sharded to take advantage of multiple cores |
| 342 | // and to allow intermediate values to reside in L3 cache. This is done by |
| 343 | // sharding the points and centers as follows: |
| 344 | // |
| 345 | // 1. Centers are sharded such that each block of centers has at most |
| 346 | // kNearestNeighborsCentersMaxBlockSize rows. |
| 347 | // 2. Points are sharded, and each block of points is multiplied with each |
nothing calls this directly
no test coverage detected