| 499 | } |
| 500 | |
| 501 | std::vector<Format> CustomOp::infer_output_format( |
| 502 | const std::vector<Format>& inputs, const Param& param) const { |
| 503 | assert_inputs_size_right(inputs); |
| 504 | for (size_t i = 0; i < inputs.size(); i++) { |
| 505 | mgb_assert( |
| 506 | inputs[i].is_default(), "the tensor format of %s:%s is not default", |
| 507 | op_type().c_str(), input_info(i).name().c_str()); |
| 508 | } |
| 509 | std::vector<Format> outputs(output_num()); |
| 510 | OpImplRef(m_impl.get())->infer_output_format_func(inputs, param, outputs); |
| 511 | |
| 512 | for (size_t i = 0; i < outputs.size(); i++) { |
| 513 | mgb_assert( |
| 514 | outputs[i].is_default(), "the tensor format of %s:%s is not default", |
| 515 | op_type().c_str(), output_info(i).name().c_str()); |
| 516 | } |
| 517 | |
| 518 | assert_outputs_size_right(outputs); |
| 519 | return outputs; |
| 520 | } |
| 521 | |
| 522 | void CustomOp::compute( |
| 523 | const std::vector<Tensor>& inputs, const Param& param, |