Should Relay infers type by default in the ExprMutator?

Development
#1

Today I found a bug in the BiasAddSimplifier pass. This pass replaces bias_add to expanding bias and add operation. However, after it modifies the bias_add expr, the subsequent exprs that rely on this expr will lose its type information. Therefore, when there’re more than one bias_add exprs, the second one will fail because it requires type information to expand the bias.

We can certainly fix this bug by some ad hoc solution. But I wonder if we should infer the type by default in Relay’s ExprMutator visitor whenever possible. I believe this can prevent many mistakes when people add new passes to Relay.

I’d like to hear thoughts from the community.

#2

This is certainly one restriction of the ExprMutator. It quite some additional effort to partially reconstruct types, mainly because there could be uncertainties in the types when it involves recursion, and backward inference.

We do need to document this problem clearly, though. The simplest way to get around this is to not rely on types of mutated input, but instead get type of the expression before mutation, for example


MyMutator::Mutate_(const Call* call) {
    // Type information is always available before mutation
    Type ret_type = call->checked_type();
    Type arg1_type = call->args[0]->checked_type();
    // Mutated result
    Expr mutated_call = IRMutator::Mutate_(call);
    // Change the result
    const Call* new_call = mutated_call.as<Call>();
    // My logics here. 
}

cc @merrymercy @jroesch