[Unity][Op] Gradient functions for high-level Relax operators

[SiriusNEO]

Intro

This PR registers gradient functions for many high-level Relax operators. Similar with Relay, the gradient function is registered as an attribute FPrimalGradient (OpAttr) of corresponding Relax operators. But the function signature is different from Relay:

using FPrimalGradient = runtime::TypedPackedFunc<tvm::Array<Expr>(
    const Var& orig_var, const Call& orig_call, const Var& output_grad, const BlockBuilder& ctx)>;

• orig_call is the orginal call expr which we want to differentiate.

• output_grad is the gradient of RHS.

• orig_var is y. It is passed to saving some calculations.

• ctx is the context which is not used right now. But we believe it is useful when it comes to dynamic shape cases and when we need to emit some bindings or do some normalizations.

For some complicate gradient functions, we introduce some high-level backward operators and put them under the namespace op.grad.xxx. All gradient functions are well tested (numerically). For more details please check Part 2 of this document.

Others

Also this PR fixes two small problems about op:

• CumsumAttrs isn't declared in the Python side.
• A small problem in the implementation about legalizing op variance.

Co-authored-by: Yixin Dong ubospica@gmail.com

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• cc @quic-sanirudh _{See #10317 for details}

_{Generated by tvm-bot}

[tqchen]

@SiriusNEO one minor note, make sure you append the co-author commit message to the end