| 46 | } |
| 47 | |
| 48 | void ReshapeGPU(OpKernelContext *context, const Tensor &input_indices_in, |
| 49 | const Tensor &input_shape_in, const Tensor &target_shape_in, |
| 50 | int output_indices_idx, int output_shape_idx) { |
| 51 | OP_REQUIRES(context, TensorShapeUtils::IsMatrix(input_indices_in.shape()), |
| 52 | errors::InvalidArgument( |
| 53 | "Input indices should be a matrix but received shape ", |
| 54 | input_indices_in.shape().DebugString())); |
| 55 | OP_REQUIRES(context, TensorShapeUtils::IsVector(input_shape_in.shape()), |
| 56 | errors::InvalidArgument( |
| 57 | "Input shape should be a vector but received shape ", |
| 58 | input_shape_in.shape().DebugString())); |
| 59 | OP_REQUIRES(context, TensorShapeUtils::IsVector(target_shape_in.shape()), |
| 60 | errors::InvalidArgument( |
| 61 | "Target shape should be a vector but received shape ", |
| 62 | target_shape_in.shape().DebugString())); |
| 63 | |
| 64 | const int64 input_rank = input_shape_in.NumElements(); |
| 65 | const int64 output_rank = target_shape_in.NumElements(); |
| 66 | const TensorShape input_shape(input_shape_in.vec<int64>()); |
| 67 | const int64 dense_size = input_shape.num_elements(); |
| 68 | const int64 nnz = input_indices_in.shape().dim_size(0); |
| 69 | |
| 70 | // Compute the output shape. Determine product of specified dimensions, and |
| 71 | // find the index of the unspecified one. |
| 72 | TensorShape output_shape; |
| 73 | int64 product = 1; |
| 74 | int unknown_index = -1; |
| 75 | auto target_shape = target_shape_in.vec<int64>(); |
| 76 | for (int d = 0; d < output_rank; ++d) { |
| 77 | const int64 size = target_shape(d); |
| 78 | if (size == -1) { |
| 79 | OP_REQUIRES( |
| 80 | context, unknown_index == -1, |
| 81 | errors::InvalidArgument("only one output dimension may be -1, " |
| 82 | "not both ", |
| 83 | unknown_index, " and ", d)); |
| 84 | unknown_index = d; |
| 85 | output_shape.AddDim(1); |
| 86 | } else { |
| 87 | OP_REQUIRES(context, size >= 0, |
| 88 | errors::InvalidArgument("size ", d, |
| 89 | " must be non-negative, not ", size)); |
| 90 | product *= size; |
| 91 | output_shape.AddDim(size); |
| 92 | } |
| 93 | } |
| 94 | if (unknown_index != -1) { |
| 95 | OP_REQUIRES( |
| 96 | context, product > 0, |
| 97 | errors::InvalidArgument("reshape cannot infer the missing " |
| 98 | "input size for an empty tensor unless all " |
| 99 | "specified input sizes are non-zero")); |
| 100 | const int64 missing = dense_size / product; |
| 101 | OP_REQUIRES( |
| 102 | context, product * missing == dense_size, |
| 103 | errors::InvalidArgument( |
| 104 | "Inferenced element num of output SparseTensor " |
| 105 | "is different from in SparseTensor. " |
no test coverage detected