Stores Layout(axis set and order) and value for dimensions.
| 84 | |
| 85 | // Stores Layout(axis set and order) and value for dimensions. |
| 86 | struct Shape { |
| 87 | Shape() : layout(Layout::UNKNOWN), dimensions() {} |
| 88 | |
| 89 | explicit Shape(Layout t) : layout(t), dimensions(Size(t)) {} |
| 90 | |
| 91 | Shape(Layout t, std::vector<int32_t> d) |
| 92 | : layout(t), dimensions(std::move(d)) {} |
| 93 | |
| 94 | bool operator==(const Shape& other) const { |
| 95 | return (layout == other.layout) && (dimensions == other.dimensions); |
| 96 | } |
| 97 | |
| 98 | bool operator!=(const Shape& other) const { return !operator==(other); } |
| 99 | |
| 100 | // All methods below are matching same methods defined in StrongShape to |
| 101 | // make sure generic algorithms work both ways. |
| 102 | |
| 103 | // Returns back a dimension or -1 if it is not found. |
| 104 | template <Axis D> |
| 105 | int32_t get() const; |
| 106 | int32_t get(Axis d) const; |
| 107 | |
| 108 | template <Axis D> |
| 109 | bool set(int32_t t); |
| 110 | bool set(Axis d, int32_t t); |
| 111 | |
| 112 | Axis axis(int index) const { return GetAxis(layout, index); } |
| 113 | |
| 114 | int index(Axis d) const { return GetAxisIndex(layout, d); } |
| 115 | |
| 116 | int64_t DimensionsProduct() const { |
| 117 | return std::accumulate(dimensions.begin(), dimensions.end(), 1ll, |
| 118 | std::multiplies<int64_t>()); |
| 119 | } |
| 120 | |
| 121 | Layout layout = Layout::UNKNOWN; |
| 122 | |
| 123 | std::vector<int32_t> dimensions; |
| 124 | }; |
| 125 | |
| 126 | std::string ToString(const Shape& s); |
| 127 |