diff --git a/include/tvm/auto_scheduler/loop_state.h b/include/tvm/auto_scheduler/loop_state.h index 489d058805c1..02392ab7fb44 100644 --- a/include/tvm/auto_scheduler/loop_state.h +++ b/include/tvm/auto_scheduler/loop_state.h @@ -429,12 +429,11 @@ class State : public ObjectRef { /*! * \brief Schedule primitive corresponds to `te::Schedule::rfactor`. * \param stage_id The index of the iterator to be factored. - * \param iter_id The iterator to be factored. + * \param it The iterator to be factored. * \param factor_iter_id The position where the new iterator is placed. * \param dag The original ComputeDAG of this state. - * \note Rfactor step will add an extra stage to the original ComputeDAG (see - * `te::Schedule::rfactor` for more details), a up-to-date ComputeDAG is stored in State's - * `current_compute_dag`. + * \note Rfactor step will add an extra stage to the original ComputeDAG (in the front of the + * target stage), a up-to-date ComputeDAG is stored in State's `current_compute_dag`. */ TVM_DLL int rfactor(int stage_id, const Iterator& it, int factor_iter_id, const ComputeDAG& dag); diff --git a/python/tvm/auto_scheduler/loop_state.py b/python/tvm/auto_scheduler/loop_state.py index 291f5f0ecafa..106cf6def87f 100644 --- a/python/tvm/auto_scheduler/loop_state.py +++ b/python/tvm/auto_scheduler/loop_state.py @@ -487,8 +487,8 @@ def rfactor(self, stage, iterator, factor_iter_id): Notes ----- - Rfactor step will insert an extra stage to the original ComputeDAG (see - `te.Schedule.rfactor` for more details). + Rfactor step will insert an extra stage to the original ComputeDAG (in the front of the + target stage). """ self.state_object, new_stage_id = _ffi_api.StateRfactor(self.state_object, self._resolve_stage_id(stage),