[Fix][Training] Fix Bind TE Grad and Handle ExternOp in Gradient Pass

mlc-ai / relax

Apache License 2.0

149 stars 75 forks source link

[Fix][Training] Fix Bind TE Grad and Handle ExternOp in Gradient Pass #207

Closed SiriusNEO closed 1 year ago

SiriusNEO commented 1 year ago

This PR contains two quick fixes for current gradient mechanism. First, it fixes the type of the attribute te_grad_bind_handler (use a cast from dict to tvm Map). Second, we ignore the ExternOp when VisitingBinding in the Gradient pass.