| 158 | |
| 159 | template <> |
| 160 | KernelGenRet opr_fill_attr<megdnn::IndexingMultiAxisVec>( |
| 161 | std::unordered_map<std::string, CCAttr>& attr_map, |
| 162 | megdnn::IndexingMultiAxisVec* opr, const TensorNDArray& tensors, |
| 163 | KernelGen::Arch arch, |
| 164 | const std::unordered_map<std::string, CCAttr>& proxy_attr) { |
| 165 | auto axis = proxy_attr.at("axis").AsOperand().shape; |
| 166 | for (size_t i = 0; i < axis.size(); ++i) { |
| 167 | attr_map["axis:" + std::to_string(i)] = CCAttr((int)axis[i]); |
| 168 | } |
| 169 | return KernelGen::KernelPack::GetKernel(KernType::IndexingMultiAxisKernel, arch); |
| 170 | } |
| 171 | |
| 172 | template <> |
| 173 | KernelGenRet opr_fill_attr<megdnn::IndexingOneHot>( |