| 185 | } |
| 186 | |
| 187 | ValueRefList subtensor_rule( |
| 188 | const Subtensor& subtensor, Span<ValueRef> inputs, Span<bool> inputs_mask, |
| 189 | const Type<ScalarValue>& scalar_type) { |
| 190 | mgb_assert(inputs.size() >= 1); |
| 191 | auto input = inputs[0]; |
| 192 | bool is_scalar; |
| 193 | mgb_assert(!inputs_mask[0], "subtensor shouldn't have scalar input"); |
| 194 | if (auto shape = input.shape()) { |
| 195 | size_t ndim = shape->ndim; |
| 196 | for (auto&& [axis, begin, end, step, idx] : subtensor.items) { |
| 197 | if (idx) { |
| 198 | ndim--; |
| 199 | } |
| 200 | } |
| 201 | is_scalar = ndim == 0; |
| 202 | } else { |
| 203 | // assume not scalar |
| 204 | is_scalar = false; |
| 205 | } |
| 206 | auto outputs = imperative::apply(subtensor, inputs); |
| 207 | if (is_scalar) { |
| 208 | outputs[0] = scalar_type.make(outputs[0]); |
| 209 | } |
| 210 | return outputs; |
| 211 | } |
| 212 | |
| 213 | ValueRefList get_var_shape_rule( |
| 214 | const GetVarShape& get_var_shape, Span<ValueRef> inputs, Span<bool> inputs_mask, |