| 66 | } |
| 67 | |
| 68 | SmallVector<TensorPtr> apply_on_physical_tensor( |
| 69 | const OpDef& def, SmallVector<TensorPtr> inputs, |
| 70 | SmallVector<LogicalTensorDesc>& output_descs, const bool& validated) { |
| 71 | auto&& op = def.cast_final_safe<IndexingOneHot>(); |
| 72 | auto&& inp = inputs[0]; |
| 73 | auto&& index = inputs[1]; |
| 74 | auto&& layout = inp->layout(); |
| 75 | auto&& index_layout = index->layout(); |
| 76 | int real_axis = static_cast<int>(op.axis); |
| 77 | if (real_axis < 0) { |
| 78 | real_axis += static_cast<int>(layout.ndim); |
| 79 | } |
| 80 | mgb_assert( |
| 81 | 0 <= real_axis && real_axis < static_cast<int>(layout.ndim), |
| 82 | "Dimension out of range (expected to be in range of [%d, %d], but got %d)", |
| 83 | 0, static_cast<int>(layout.ndim) - 1, op.axis); |
| 84 | DnnOprCaller<megdnn::IndexingOneHot> dnn_op(inp->comp_node(), real_axis); |
| 85 | auto tlayout = dnn_op.deduce_layout(layout, index_layout); |
| 86 | auto out = Tensor::make(tlayout, inp->comp_node()); |
| 87 | if (layout.is_empty()) { |
| 88 | return {out}; |
| 89 | } |
| 90 | dnn_op.exec_with_ws(inp, index, out); |
| 91 | return {out}; |
| 92 | } |
| 93 | |
| 94 | OP_TRAIT_REG(IndexingOneHot, IndexingOneHot) |
| 95 | .infer_output_attrs_fallible(infer_output_attrs_fallible) |
nothing calls this directly
no test coverage detected