| 33 | namespace tensorflow { |
| 34 | |
| 35 | void Reshape(OpKernelContext *context, const Tensor &input_indices_in, |
| 36 | const Tensor &input_shape_in, const Tensor &target_shape_in, |
| 37 | int output_indices_idx, int output_shape_idx) { |
| 38 | OP_REQUIRES(context, TensorShapeUtils::IsMatrix(input_indices_in.shape()), |
| 39 | errors::InvalidArgument( |
| 40 | "Input indices should be a matrix but received shape ", |
| 41 | input_indices_in.shape().DebugString())); |
| 42 | OP_REQUIRES(context, TensorShapeUtils::IsVector(input_shape_in.shape()), |
| 43 | errors::InvalidArgument( |
| 44 | "Input shape should be a vector but received shape ", |
| 45 | input_shape_in.shape().DebugString())); |
| 46 | OP_REQUIRES(context, TensorShapeUtils::IsVector(target_shape_in.shape()), |
| 47 | errors::InvalidArgument( |
| 48 | "Target shape should be a vector but received shape ", |
| 49 | target_shape_in.shape().DebugString())); |
| 50 | |
| 51 | const int64 input_rank = input_shape_in.NumElements(); |
| 52 | const int64 output_rank = target_shape_in.NumElements(); |
| 53 | TensorShape input_shape; |
| 54 | OP_REQUIRES_OK(context, TensorShape::BuildTensorShape( |
| 55 | input_shape_in.vec<int64>(), &input_shape)); |
| 56 | const int64 dense_size = input_shape.num_elements(); |
| 57 | const int64 nnz = input_indices_in.shape().dim_size(0); |
| 58 | |
| 59 | // Compute the output shape. Determine product of specified dimensions, and |
| 60 | // find the index of the unspecified one. |
| 61 | TensorShape output_shape; |
| 62 | int64 product = 1; |
| 63 | int unknown_index = -1; |
| 64 | auto target_shape = target_shape_in.vec<int64>(); |
| 65 | for (int d = 0; d < output_rank; ++d) { |
| 66 | const int64 size = target_shape(d); |
| 67 | if (size == -1) { |
| 68 | OP_REQUIRES( |
| 69 | context, unknown_index == -1, |
| 70 | errors::InvalidArgument("only one output dimension may be -1, " |
| 71 | "not both ", |
| 72 | unknown_index, " and ", d)); |
| 73 | unknown_index = d; |
| 74 | output_shape.AddDim(1); |
| 75 | } else { |
| 76 | OP_REQUIRES(context, size >= 0, |
| 77 | errors::InvalidArgument("size ", d, |
| 78 | " must be non-negative, not ", size)); |
| 79 | product *= size; |
| 80 | output_shape.AddDim(size); |
| 81 | } |
| 82 | } |
| 83 | if (unknown_index != -1) { |
| 84 | OP_REQUIRES( |
| 85 | context, product > 0, |
| 86 | errors::InvalidArgument("reshape cannot infer the missing " |
| 87 | "input size for an empty tensor unless all " |
| 88 | "specified input sizes are non-zero")); |
| 89 | const int64 missing = dense_size / product; |
| 90 | OP_REQUIRES( |
| 91 | context, product * missing == dense_size, |
| 92 | errors::InvalidArgument( |