[TE] reverse-mode autodiff without any optimization

[yzhliu]

This is the first PR to bring in previously-implemented tensor-level autodiff.

This PR does not include any optimization, thus produces bad performance. Will submit optimization pass in another two or three PRs, so that not to put too much pressure on reviewers.

Also credit to @sgrechanik-h as I mentioned in the header of each file.

RFC: https://discuss.tvm.ai/t/rfc-bring-in-tensor-expression-autodiff

Please help to review @sgrechanik-h @MarisaKirisame @junrushao1994 @tqchen @hzfan

[MarisaKirisame]

can you switch the order? most ad code use original result first, derivatives later.

[yzhliu]

Looking into a bit more, the order actually makes difference. When original init value is different from its derivative init value, and they depends on each other during calculation, we must calculate derivative first (using origin's init value), switch the order in tvm makes the origin value be replaced before using, produces incorrect results.

One example is in the test case,

 def fcombine(x, y):
        return x*y

    def fidentity(t0):
        return tvm.tir.const(1, t0)

    prod = te.comm_reducer(fcombine, fidentity, name='prod')
    B = te.compute((10, 10), lambda i, j: prod(A0[i, k] + A0[k, i], axis=k), name='B')
    check_grad(B, A0)

Correct result (derivative first):

produce B.jacobian {
  for (i, 0, 10) {
    for (j, 0, 10) {
      for (jac_i0, 0, 10) {
        for (jac_i1, 0, 10) {
          B.jacobian.v0[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = 0f
          B.jacobian.v1[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = 1f
          for (k, 0, 10) {
            B.jacobian.v0[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = ((B.jacobian.v0[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)]*(A0[((i*10) + k)] + A0[((k*10) + i)])) + ((float32(((jac_i0 == i) && (jac_i1 == k))) + float32(((jac_i0 == k) && (jac_i1 == i))))*B.jacobian.v1[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)]))
            B.jacobian.v1[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = (B.jacobian.v1[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)]*(A0[((i*10) + k)] + A0[((k*10) + i)]))
          }
        }
      }
    }
  }
}
Output B.jacobian.v0

Incorrect result (origin first):

produce B.jacobian {
  for (i, 0, 10) {
    for (j, 0, 10) {
      for (jac_i0, 0, 10) {
        for (jac_i1, 0, 10) {
          B.jacobian.v0[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = 1f
          B.jacobian.v1[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = 0f
          for (k, 0, 10) {
            B.jacobian.v0[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = (B.jacobian.v0[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)]*(A0[((i*10) + k)] + A0[((k*10) + i)]))
            B.jacobian.v1[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = ((B.jacobian.v1[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)]*(A0[((i*10) + k)] + A0[((k*10) + i)])) + ((float32(((jac_i0 == i) && (jac_i1 == k))) + float32(((jac_i0 == k) && (jac_i1 == i))))*B.jacobian.v0[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)]))
          }
        }
      }
    }
  }
}
Output B.jacobian.v1

[sergei-grechanik]

Looks more like a bug in lowering of tupled reductions rather than an intended behavior, might deserve a separate bug report.

[MarisaKirisame]

@tqchen can you take a quick look and see if it is a bug in tuple reductions?

[MarisaKirisame]

Is it possible to have reduce inside other expression as well?

[yzhliu]

seems to be a bit difficult. do you have an concrete example in mind?

[MarisaKirisame]

There are no concrete example. If you dont think it will happend then leave it as is.

[hzfan]

So the default behavior is to return a Jacobian instead of adjoint, right?

[yzhliu]

that's right. more precisely, that's because the arguments are one output and multiple inputs, instead of one input and multiple outputs. if y is the only output, dy/dx is jacobian, it's also adjoint(x) for the previous layer. it depends on what aspect you want to emphasize, you use different terms.

[sgrechanik-h]

Thanks for reviving this. The PR looks good to me, but I'm obviously partial.

[yzhliu]

@MarisaKirisame @tqchen @hzfan Could you review again?

[MarisaKirisame]

@yzhliu can you do forward mode automatic differentiation? It is easy considering you have jacobian - you only need to do JacobianVectorProduct instead of VectorJacbianProduct

It is useful in higher order derivative a la hessian vector product.

[MarisaKirisame]

(of course, not in this PR)

[yzhliu]

@MarisaKirisame sure, I will try.

[yzhliu]

CI's green. @tqchen @hzfan check if there's anything needs to be addressed.

[yzhliu]

Kindly ping @tqchen , can we merge if it looks good?

[tqchen]

Thanks @yzhliu @sgrechanik-h @MarisaKirisame @hzfan