* @brief Prepare TensorListShape based on `in_shape` number of channels and spatial output * dimensions provided in `output_sizes`. */
| 143 | * dimensions provided in `output_sizes`. |
| 144 | */ |
| 145 | TensorListShape<tensor_ndim> GetOutputShape( |
| 146 | const TensorListShape<tensor_ndim> &in_shape, |
| 147 | span<const TensorShape<ndim>> output_sizes) { |
| 148 | assert(in_shape.num_samples() == static_cast<int>(output_sizes.size())); |
| 149 | TensorListShape<tensor_ndim> shape; |
| 150 | shape.resize(in_shape.num_samples(), tensor_ndim); |
| 151 | for (int i = 0; i < in_shape.num_samples(); i++) { |
| 152 | auto out_tshape = shape.tensor_shape_span(i); |
| 153 | int in_d = 0; |
| 154 | for (int j = 0; j < tensor_ndim; j++) { |
| 155 | out_tshape[j] = (j == channel_dim) |
| 156 | ? in_shape.tensor_shape_span(i)[channel_dim] |
| 157 | : output_sizes[i][in_d++]; |
| 158 | } |
| 159 | } |
| 160 | return shape; |
| 161 | } |
| 162 | |
| 163 | /** |
| 164 | * @brief Check if `out_shape` matches the number of channels of `in_shape` and spatial dimensions |