Make a sparse spec into a dense index spec. The sparse spec does not correspond to the number of dimensions Make a dense spec that corresponds to the number of dimensions For example suppose foo[...,3:, 2] on foo.shape=(2,2,3,4) then we need to produce the missing begin_mask, end_mask for the first two dimensions i.e. foo[:, :, 3:, 2].
| 2639 | // we need to produce the missing begin_mask, end_mask for the first two |
| 2640 | // dimensions i.e. foo[:, :, 3:, 2]. |
| 2641 | static LogicalResult BuildDenseSliceSpec(const SparseSliceSpec &sparse, |
| 2642 | DenseSliceSpec *dense) { |
| 2643 | // Build expanded dense begin, end, strides, begin_mask, end_mask, and |
| 2644 | // shrink_axis_mask. |
| 2645 | dense->begin.resize(dense->dims); |
| 2646 | dense->end.resize(dense->dims); |
| 2647 | dense->strides.resize(dense->dims); |
| 2648 | dense->begin_mask = 0; |
| 2649 | dense->end_mask = 0; |
| 2650 | dense->shrink_axis_mask = 0; |
| 2651 | |
| 2652 | // Count number of new_axis after ellipsis. This helps in calculating the |
| 2653 | // number of dimensions ellipsis represents in the sparse spec. |
| 2654 | bool ellipsis_seen = false; |
| 2655 | int num_new_axis_after_ellipsis = 0; |
| 2656 | for (int sparse_index = 0; sparse_index < sparse.dims; ++sparse_index) { |
| 2657 | if (ellipsis_seen && IsSet(sparse.new_axis_mask, sparse_index)) |
| 2658 | num_new_axis_after_ellipsis++; |
| 2659 | if (IsSet(sparse.ellipsis_mask, sparse_index)) ellipsis_seen = true; |
| 2660 | } |
| 2661 | |
| 2662 | int dense_index = 0; |
| 2663 | for (int sparse_index = 0; sparse_index < sparse.dims; ++sparse_index) { |
| 2664 | if (IsSet(sparse.new_axis_mask, sparse_index)) continue; |
| 2665 | if (IsSet(sparse.ellipsis_mask, sparse_index)) { |
| 2666 | auto next_index = std::min(dense->dims - (sparse.dims - sparse_index) + |
| 2667 | 1 + num_new_axis_after_ellipsis, |
| 2668 | dense->dims); |
| 2669 | // Expand ellipsis into the appropriate dense indices. From current index |
| 2670 | // until next_index, all dimensions would have begin and end masks set and |
| 2671 | // stride 1, i.e., get all elements in those dimensions. |
| 2672 | for (; dense_index < next_index; ++dense_index) { |
| 2673 | dense->begin[dense_index] = dense->end[dense_index] = 0; |
| 2674 | dense->strides[dense_index] = 1; |
| 2675 | Set(dense->begin_mask, dense_index); |
| 2676 | Set(dense->end_mask, dense_index); |
| 2677 | } |
| 2678 | continue; |
| 2679 | } |
| 2680 | assert(dense_index < dense->dims); |
| 2681 | // Copy over the sparse indices to dense indices if ellipsis_mask and |
| 2682 | // new_axis_mask are not set. |
| 2683 | dense->begin[dense_index] = sparse.begin[sparse_index]; |
| 2684 | dense->end[dense_index] = sparse.end[sparse_index]; |
| 2685 | dense->strides[dense_index] = sparse.strides[sparse_index]; |
| 2686 | CopyBit(sparse.begin_mask, sparse_index, dense->begin_mask, dense_index); |
| 2687 | CopyBit(sparse.end_mask, sparse_index, dense->end_mask, dense_index); |
| 2688 | CopyBit(sparse.shrink_axis_mask, sparse_index, dense->shrink_axis_mask, |
| 2689 | dense_index); |
| 2690 | dense_index++; |
| 2691 | } |
| 2692 | return success(); |
| 2693 | } |
| 2694 | |
| 2695 | // For the given `input_shape`, calculates the sliced shape using the given |
| 2696 | // `begin`, `end`, and `stride` ranges and `begin_mask`, `end_mask`, and |