* @brief Simplifies a shape by collapsing dimensions that are not strided * * @param[in] rank number of dimensions * @param[in] shape * @param[in] stride * @param[out] out_shape * @param[out] out_strides * @return int out_rank */
| 199 | * @return int out_rank |
| 200 | */ |
| 201 | static int Simplify(int rank, int64_t *shape, int64_t *stride, int64_t *out_shape, int64_t *out_strides) |
| 202 | { |
| 203 | if (rank <= 1) // Nothing to simplify |
| 204 | { |
| 205 | if (rank == 1) |
| 206 | { |
| 207 | out_shape[0] = shape[0]; |
| 208 | out_strides[0] = stride[0]; |
| 209 | } |
| 210 | return rank; |
| 211 | } |
| 212 | |
| 213 | int out_rank = 0; |
| 214 | int64_t vol = shape[0]; |
| 215 | for (int d = 1; d < rank; d++) |
| 216 | { |
| 217 | if (stride[d - 1] != shape[d] * stride[d]) |
| 218 | { |
| 219 | out_strides[out_rank] = stride[d - 1]; |
| 220 | out_shape[out_rank] = vol; |
| 221 | vol = shape[d]; |
| 222 | out_rank++; |
| 223 | } |
| 224 | else |
| 225 | { |
| 226 | vol *= shape[d]; |
| 227 | } |
| 228 | } |
| 229 | out_strides[out_rank] = stride[rank - 1]; |
| 230 | out_shape[out_rank] = vol; |
| 231 | out_rank++; |
| 232 | return out_rank; |
| 233 | } |
| 234 | |
| 235 | /** |
| 236 | * @brief Reshapes a simplified shape (non-strided dimensions are collapsed) to a target shape if possible. |