| 153 | namespace stack { |
| 154 | |
| 155 | TensorLayout stack_layout_deduce( |
| 156 | const SmallVector<const TensorLayout*> inputs, int axis) { |
| 157 | size_t nr_inp = inputs.size(); |
| 158 | auto&& inp_layout0 = *inputs[0]; |
| 159 | for (size_t i = 1; i < nr_inp; ++i) { |
| 160 | mgb_assert( |
| 161 | inp_layout0.eq_shape(*inputs[i]), |
| 162 | "Stack input shape mismatch: %s vs %s", inp_layout0.to_string().c_str(), |
| 163 | inputs[i]->to_string().c_str()); |
| 164 | } |
| 165 | |
| 166 | TensorLayout oup_layout{TensorShape{inp_layout0}, inp_layout0.dtype}; |
| 167 | oup_layout.add_axis_cont_inplace(axis); |
| 168 | oup_layout.shape[axis] = nr_inp; |
| 169 | oup_layout.init_contiguous_stride(); |
| 170 | return oup_layout; |
| 171 | } |
| 172 | |
| 173 | auto apply_on_var_node(const OpDef& def, const VarNodeArray& inputs) { |
| 174 | auto&& op = static_cast<const Stack&>(def); |
no test coverage detected