* refactor general_grad and fix some bugs

* add TODO: support prune logic deeper