| 30 | } |
| 31 | |
| 32 | void Compute(OpKernelContext* context) override { |
| 33 | const Tensor& input_axis = context->input(0); |
| 34 | const Tensor& input_indices = context->input(1); |
| 35 | const Tensor& input_values = context->input(2); |
| 36 | const Tensor& input_shape = context->input(3); |
| 37 | |
| 38 | OP_REQUIRES(context, TensorShapeUtils::IsScalar(input_axis.shape()), |
| 39 | errors::InvalidArgument( |
| 40 | "Input axis should be a scalar but received shape ", |
| 41 | input_axis.shape().DebugString())); |
| 42 | OP_REQUIRES(context, TensorShapeUtils::IsMatrix(input_indices.shape()), |
| 43 | errors::InvalidArgument( |
| 44 | "Input indices should be a matrix but received shape ", |
| 45 | input_indices.shape().DebugString())); |
| 46 | OP_REQUIRES(context, TensorShapeUtils::IsVector(input_values.shape()), |
| 47 | errors::InvalidArgument( |
| 48 | "Input values should be a vector but received shape ", |
| 49 | input_indices.shape().DebugString())); |
| 50 | OP_REQUIRES(context, TensorShapeUtils::IsVector(input_shape.shape()), |
| 51 | errors::InvalidArgument( |
| 52 | "Input shape should be a vector but received shape ", |
| 53 | input_shape.shape().DebugString())); |
| 54 | |
| 55 | const int64 split_dim = input_axis.scalar<int64>()(); |
| 56 | |
| 57 | OP_REQUIRES( |
| 58 | context, |
| 59 | input_shape.dim_size(0) && split_dim < input_shape.vec<int64>().size(), |
| 60 | errors::InvalidArgument( |
| 61 | "Input split_dim should be between 0 and rank (", |
| 62 | input_shape.vec<int64>().size(), "), got ", split_dim)); |
| 63 | |
| 64 | OP_REQUIRES( |
| 65 | context, |
| 66 | num_split_ >= 1 && num_split_ <= input_shape.vec<int64>()(split_dim), |
| 67 | errors::InvalidArgument("Input num_split should be between 1 " |
| 68 | "and the splitting dimension size (", |
| 69 | input_shape.vec<int64>()(split_dim), "), got ", |
| 70 | num_split_)); |
| 71 | |
| 72 | sparse::SparseTensor sparse_tensor; |
| 73 | OP_REQUIRES_OK(context, |
| 74 | sparse::SparseTensor::Create( |
| 75 | input_indices, input_values, |
| 76 | TensorShape(input_shape.vec<int64>()), &sparse_tensor)); |
| 77 | |
| 78 | std::vector<sparse::SparseTensor> outputs; |
| 79 | OP_REQUIRES_OK(context, |
| 80 | sparse::SparseTensor::Split<T>(sparse_tensor, split_dim, |
| 81 | num_split_, &outputs)); |
| 82 | |
| 83 | for (int slice_index = 0; slice_index < num_split_; ++slice_index) { |
| 84 | context->set_output(slice_index, outputs[slice_index].indices()); |
| 85 | context->set_output(slice_index + num_split_, |
| 86 | outputs[slice_index].values()); |
| 87 | Tensor* shape = nullptr; |
| 88 | OP_REQUIRES_OK(context, context->allocate_output( |
| 89 | slice_index + 2 * num_split_, |
nothing calls this directly
no test coverage detected