| 360 | explicit TileGradientOp(OpKernelConstruction* context) : OpKernel(context) {} |
| 361 | |
| 362 | void Compute(OpKernelContext* context) override { |
| 363 | const Tensor& input = context->input(0); |
| 364 | const Tensor& multiples = context->input(1); |
| 365 | OP_REQUIRES( |
| 366 | context, IsLegacyVector(multiples.shape()), |
| 367 | errors::InvalidArgument("Expected multiples to be 1-D, but got shape ", |
| 368 | multiples.shape().DebugString())); |
| 369 | OP_REQUIRES(context, input.dims() == multiples.NumElements(), |
| 370 | errors::InvalidArgument( |
| 371 | "Expected multiples argument to be a vector of length ", |
| 372 | input.dims(), " but got length ", multiples.dim_size(0))); |
| 373 | |
| 374 | const int input_dims = input.dims(); |
| 375 | |
| 376 | // Eigen doesn't support scalars on the GPU, so handle 0-D specially |
| 377 | if (input_dims == 0) { |
| 378 | context->set_output(0, input); |
| 379 | return; |
| 380 | } |
| 381 | |
| 382 | const gtl::ArraySlice<Tmultiples> multiples_array( |
| 383 | multiples.flat<Tmultiples>().data(), input_dims); |
| 384 | TensorShape output_shape; |
| 385 | std::vector<Tmultiples> input_dim_size_vec; |
| 386 | for (int i = 0; i < input_dims; ++i) { |
| 387 | OP_REQUIRES( |
| 388 | context, multiples_array[i] > 0, |
| 389 | errors::InvalidArgument("Expected multiples[", i, "] > 0, but got ", |
| 390 | multiples_array[i])); |
| 391 | OP_REQUIRES(context, input.dim_size(i) % multiples_array[i] == 0, |
| 392 | errors::InvalidArgument("Expected input_dim[", i, |
| 393 | "] to be divisible by multiples[", i, |
| 394 | "], but ", input.dim_size(i), " % ", |
| 395 | multiples_array[i], " != 0")); |
| 396 | output_shape.AddDim(input.dim_size(i) / multiples_array[i]); |
| 397 | input_dim_size_vec.push_back(input.dim_size(i)); |
| 398 | } |
| 399 | if (output_shape == input.shape()) { |
| 400 | context->set_output(0, input); |
| 401 | return; |
| 402 | } |
| 403 | Tensor* result = nullptr; |
| 404 | OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &result)); |
| 405 | |
| 406 | #define HANDLE_DIM(DT, NDIM) \ |
| 407 | if (context->input(0).dtype() == DT && input_dims == NDIM) { \ |
| 408 | HandleCase<DT, NDIM>(context, input_dim_size_vec, multiples_array, \ |
| 409 | result); \ |
| 410 | return; \ |
| 411 | } |
| 412 | |
| 413 | #define HANDLE_TYPE(T) \ |
| 414 | HANDLE_DIM(T, 1) \ |
| 415 | HANDLE_DIM(T, 2) \ |
| 416 | HANDLE_DIM(T, 3) \ |
| 417 | HANDLE_DIM(T, 4) \ |
| 418 | HANDLE_DIM(T, 5) \ |
| 419 | HANDLE_DIM(T, 6) \ |
nothing calls this directly
no test coverage detected