| 36 | explicit SparseReorderOp(OpKernelConstruction* context) : OpKernel(context) {} |
| 37 | |
| 38 | void Compute(OpKernelContext* context) override { |
| 39 | const Tensor& input_ind = context->input(0); |
| 40 | OP_REQUIRES(context, TensorShapeUtils::IsMatrix(input_ind.shape()), |
| 41 | errors::InvalidArgument( |
| 42 | "Input indices should be a matrix but received shape ", |
| 43 | input_ind.shape().DebugString())); |
| 44 | |
| 45 | const Tensor& input_val = context->input(1); |
| 46 | OP_REQUIRES(context, TensorShapeUtils::IsVector(input_val.shape()), |
| 47 | errors::InvalidArgument( |
| 48 | "Input values should be a vector but received shape ", |
| 49 | input_val.shape().DebugString())); |
| 50 | |
| 51 | const Tensor& input_shape_in = context->input(2); |
| 52 | OP_REQUIRES(context, TensorShapeUtils::IsVector(input_shape_in.shape()), |
| 53 | errors::InvalidArgument( |
| 54 | "Input shape should be a vector but received shape ", |
| 55 | input_shape_in.shape().DebugString())); |
| 56 | |
| 57 | const TensorShape input_shape(input_shape_in.vec<int64>()); |
| 58 | |
| 59 | gtl::InlinedVector<int64, 8> std_order(input_shape.dims()); |
| 60 | std::iota(std_order.begin(), std_order.end(), 0); |
| 61 | |
| 62 | // Check if the sparse tensor is already ordered correctly |
| 63 | sparse::SparseTensor input_sp; |
| 64 | OP_REQUIRES_OK( |
| 65 | context, sparse::SparseTensor::Create(input_ind, input_val, input_shape, |
| 66 | std_order, &input_sp)); |
| 67 | |
| 68 | if (input_sp.IndicesValid().ok()) { |
| 69 | context->set_output(0, input_sp.indices()); |
| 70 | context->set_output(1, input_sp.values()); |
| 71 | } else { |
| 72 | // Deep-copy the input Tensors, then reorder in-place |
| 73 | sparse::SparseTensor reordered_sp; |
| 74 | OP_REQUIRES_OK(context, |
| 75 | sparse::SparseTensor::Create(tensor::DeepCopy(input_ind), |
| 76 | tensor::DeepCopy(input_val), |
| 77 | input_shape, &reordered_sp)); |
| 78 | reordered_sp.Reorder<T>(std_order); |
| 79 | context->set_output(0, reordered_sp.indices()); |
| 80 | context->set_output(1, reordered_sp.values()); |
| 81 | } |
| 82 | } |
| 83 | }; |
| 84 | |
| 85 | #define REGISTER_KERNELS(type) \ |
nothing calls this directly
no test coverage detected