| 997 | } |
| 998 | |
| 999 | bool Split::infer_shape( |
| 1000 | size_t out_idx, TensorShape& dest, const cg::static_infer::InpVal& inp) { |
| 1001 | mgb_assert(inp.run_id > 0, "run id should be a positive number"); |
| 1002 | if (inp.run_id != m_output_shape_version) { |
| 1003 | std::vector<size_t> partition; |
| 1004 | auto ishp = inp.val.at(0).shape(); |
| 1005 | auto axis = m_opt.axis; |
| 1006 | if (axis < 0) |
| 1007 | axis += ishp.ndim; |
| 1008 | if (m_opt.method == Options::Method::SPECIFY) { |
| 1009 | for (size_t i = 0; i < m_opt.nr_part; ++i) { |
| 1010 | auto&& val = inp.val.at(i + 1).value(); |
| 1011 | mgb_assert(val.shape().is_scalar(), "shapes for Split must be scalars"); |
| 1012 | size_t cvt; |
| 1013 | static_cast_dtype_safe(&cvt, val.dtype(), val.raw_ptr()); |
| 1014 | partition.push_back(cvt); |
| 1015 | } |
| 1016 | } else { |
| 1017 | partition = m_opt.callback(ishp.shape[axis]); |
| 1018 | mgb_assert( |
| 1019 | partition.size() == m_opt.nr_part, |
| 1020 | "nr_part=%zu but split callback returned %zu parts", m_opt.nr_part, |
| 1021 | partition.size()); |
| 1022 | } |
| 1023 | size_t size = 0; |
| 1024 | for (size_t i = 0; i < m_opt.nr_part; ++i) { |
| 1025 | auto p = partition[i]; |
| 1026 | size += p; |
| 1027 | |
| 1028 | auto&& cur = m_output_spec[i].shape; |
| 1029 | cur = ishp; |
| 1030 | cur.shape[axis] = p; |
| 1031 | } |
| 1032 | mgb_assert( |
| 1033 | size == ishp.shape[axis], |
| 1034 | "split size sums to %zd, but shape at the axis is %zd", size, |
| 1035 | ishp.shape[axis]); |
| 1036 | m_output_shape_version = inp.run_id; |
| 1037 | } |
| 1038 | |
| 1039 | dest = m_output_spec.at(out_idx).shape; |
| 1040 | return true; |
| 1041 | } |
| 1042 | |
| 1043 | void Split::init_output_comp_node() { |
| 1044 | auto&& conf_node = config().comp_node(); |