| 55 | } |
| 56 | |
| 57 | auto apply_on_var_node(const OpDef& def, const VarNodeArray& inputs) { |
| 58 | auto&& op = def.cast_final_safe<IndexingOneHot>(); |
| 59 | mgb_assert(inputs.size() == 2); |
| 60 | int real_axis = static_cast<int>(op.axis); |
| 61 | if (real_axis < 0) { |
| 62 | real_axis += static_cast<int>(op.ndim); |
| 63 | } |
| 64 | OperatorNodeConfig config{op.make_name()}; |
| 65 | return opr::IndexingOneHot::make(inputs[0], inputs[1], real_axis, config); |
| 66 | } |
| 67 | |
| 68 | SmallVector<TensorPtr> apply_on_physical_tensor( |
| 69 | const OpDef& def, SmallVector<TensorPtr> inputs, |