I am trying to see if I can easily add the linearity support for aten.mul to allow Partial placement to propagate through. But it turns out that I have to completely rework the current linearity propagation.

In short, before this PR, linearity mainly support aten.add and some trival ops. It is done by allowing input Partial to propagate, and in the meanwhile, redistribute Replicate inputs to Partial to preserve the single device semantic, i.e suppose we want to execute aten.add(lhs, rhs) on 2 ranks:

• lhs is partial, value on rank 0: r0, lhs value on rank 1: r1
• rhs is replicate, value: a

Then in order to preserve single device semantic (which should produce the value of a + r0 + r1), we do rhs/world_size first, then add rhs to lhs. This means every operand would first need be partial, then we can add them together.

But this become non-true for multiplicative operations, like aten.mul, for aten.mul, assuming the same aten.mul(lhs, rhs) and value, we don't need to divide lhs by world_size to preserve single device semantic, b.c. a* (r0+r1) = a* r0 + a* r1

So to accomodate the difference of add/mul, in this PR I:

• change linearity to be a int to support different linearity types, add linearity and multiplicative are separate
• add checks to ensure only a subset of partial types can support linearity (namely partial-sum/avg)
• handle the linearity type plumbing through the pointwise ops.
• add mul.Tensor/Scalar to be the multiplicative linearity
• added the tests to show that the partial placements can be propagated with aten.mul

Fixes #ISSUE_NUMBER

cc @H-Huang @awgu @fegin @fduwjj @wz337 @wconstab @d4l3k