Compute common reduce parameters that'll be used for SparseTensor reductions. Usage: ReduceDetails reduction = SparseTensorReduceHelper(sp, axes, keep_dims); sp.Reorder(reduction.reorder_dims); for (const auto& g : sp.group(reduction.group_by_dims)) { ... } // Set output shape to reduction.reduced_shape.
| 53 | // } |
| 54 | // // Set output shape to reduction.reduced_shape. |
| 55 | ReduceDetails SparseTensorReduceHelper(const SparseTensor &sp, |
| 56 | gtl::ArraySlice<int32> axes_slice, |
| 57 | bool keep_dims) { |
| 58 | ReduceDetails reduction; |
| 59 | |
| 60 | std::vector<int32> reduction_axes(axes_slice.begin(), axes_slice.end()); |
| 61 | int ndims = sp.dims(); |
| 62 | for (int64 i = 0; i < reduction_axes.size(); ++i) { |
| 63 | reduction_axes[i] = (reduction_axes[i] + ndims) % ndims; |
| 64 | } |
| 65 | std::sort(reduction_axes.begin(), reduction_axes.end()); |
| 66 | |
| 67 | // (0) Calculate the grouping dimensions: |
| 68 | // group_by_dims == {0, .., NDIMS-1} \ reduction_axes. |
| 69 | std::vector<int64> perm(ndims); |
| 70 | std::iota(perm.begin(), perm.end(), 0); |
| 71 | |
| 72 | // Requires perm and reduction_axes_ be sorted; group_by_dims will be |
| 73 | // sorted as well. |
| 74 | std::set_difference( |
| 75 | perm.begin(), perm.end(), reduction_axes.begin(), reduction_axes.end(), |
| 76 | std::inserter(reduction.group_by_dims, reduction.group_by_dims.begin())); |
| 77 | |
| 78 | // Now append the rest of the axes (the complement of group_by_dims_); |
| 79 | // result is used by Reorder(). |
| 80 | reduction.reorder_dims = reduction.group_by_dims; |
| 81 | std::set_difference(perm.begin(), perm.end(), reduction.group_by_dims.begin(), |
| 82 | reduction.group_by_dims.end(), |
| 83 | std::back_inserter(reduction.reorder_dims)); |
| 84 | |
| 85 | // (1) Calculate the shape after reduction. |
| 86 | auto sp_shape = sp.shape(); |
| 87 | std::vector<int64> out_dim_sizes; |
| 88 | if (keep_dims) { |
| 89 | out_dim_sizes.reserve(ndims); |
| 90 | auto beg = reduction.group_by_dims.begin(); |
| 91 | auto end = reduction.group_by_dims.end(); |
| 92 | for (int d = 0; d < ndims; ++d) { |
| 93 | if (std::find(beg, end, d) == end) { |
| 94 | out_dim_sizes.push_back(1); // A reduced axis. |
| 95 | } else { |
| 96 | out_dim_sizes.push_back(sp_shape[d]); |
| 97 | } |
| 98 | } |
| 99 | } else { |
| 100 | out_dim_sizes = sp.PickDims(reduction.group_by_dims); |
| 101 | } |
| 102 | |
| 103 | reduction.reduced_shape = TensorShape(out_dim_sizes); |
| 104 | return reduction; |
| 105 | } |
| 106 | |
| 107 | Status ValidateInputs(const Tensor *shape_t, const Tensor *reduction_axes_t) { |
| 108 | // indices and values are validated in SparseTensor ctor. |