| 451 | } |
| 452 | |
| 453 | std::vector<Shape> CustomOp::infer_output_shape( |
| 454 | const std::vector<Shape>& inputs, const Param& param) const { |
| 455 | assert_inputs_size_right(inputs); |
| 456 | for (size_t i = 0; i < inputs_info().size(); i++) { |
| 457 | assert_arg_shape_dim_right(inputs[i], input_info(i)); |
| 458 | } |
| 459 | |
| 460 | std::vector<Shape> outputs(output_num()); |
| 461 | OpImplRef(m_impl.get())->infer_output_shape_func(inputs, param, outputs); |
| 462 | for (size_t i = 0; i < outputs_info().size(); i++) { |
| 463 | assert_arg_shape_dim_right(outputs[i], output_info(i)); |
| 464 | } |
| 465 | |
| 466 | assert_outputs_size_right(outputs); |
| 467 | return outputs; |
| 468 | } |
| 469 | |
| 470 | std::vector<DType> CustomOp::infer_output_dtype( |
| 471 | const std::vector<DType>& inputs, const Param& param) const { |