[Sema] [AutoDiff] Derive @_fieldwiseProductSpace when tangent/cotangent is Self

[rxwei]

In Differentiable derived conformances, fix a bug where @_fieldwiseProductSpace is not added to derived typealiases when TangentVector and/or CotangentVector are inferred as Self. Without this change, field access x.a could not be differentiated.

struct AdditiveTangentIsSelf : AdditiveArithmetic, Differentiable {
  var a: Float
}
let _: @autodiff (AdditiveTangentIsSelf) -> Float = { x in
  x.a + x.a
}

[rxwei]

@swift-ci please test tensorflow

[rxwei]

@swift-ci please test tensorflow

[rxwei]

@swift-ci please test tensorflow

[rxwei]

@swift-ci please test tensorflow

[rxwei]

@swift-ci please test tensorflow

[rxwei]

@swift-ci please test tensorflow

[dan-zheng]

@swift-ci Please test tensorflow

[dan-zheng]

@swift-ci Please test tensorflow

[dan-zheng]

@swift-ci Please test tensorflow

[rxwei]

@swift-ci please test tensorflow

[rxwei]

@swift-ci please test tensorflow macOS

[dan-zheng]

Good catch!

[dan-zheng]

Do you want to change this to a note?

[rxwei]

No, here's why. The DifferentiationInvoker infra keeps track of how each differentiation task is created. Only the diagnostic on the very bottom of the invocation stack should be an error.

func f0(x) {
  f1()
}
func f1(x) {
  f2(x)
}
func f2(x) {
  f3(x) // non-differentiable
}
derivative(of: f0)

If expression f3() is not differentiable, we should not emit an error at f3() because it's not the place where the user requested the derivative and the user should not be told to change the code here before knowing what's actually causing differentiation to occur. Instead, we emit notes along the invocation stack until we reach the user's differentiation request, which is always the bottom of the invocation stack.

func f0(x) {
  f1() // note: when differentiating this function call
}
func f1(x) {
  f2(x) // note: when differentiating this function call
}
func f2(x) {
  f3(x) // note: expression is not differentiable
}
derivative(of: f0) // error: f0 is not differentiable

As to why this cannot be changed to a note today: We used to have DifferentialOperator as a case in DifferentiationInvoker, and the differential operator was, as you know, #gradient. When we switched over to generalized differentiability™, there is no more compiler-known differential operator -- function conversion to @autodiff is the only formal way to trigger differentiation in expressions. We didn't add a case for function conversion in DifferentiationInvoker, so the user request is not recorded, and there is no error at the user request. If we change this error to a note today, the pass pipeline would assume there's no error, keep running other passes on partially/incorrectly transformed code, and cause crashers or undefined runtime behavior.

[rxwei]

PS: This definitely needs to be fixed. Let me know if you wanna make diagnostics better :)

[dan-zheng]

Thanks for the detailed info! Definitely interested in improving AD diagnostics sometime.

[rxwei]

@swift-ci please test tensorflow

[rxwei]

@swift-ci please test tensorflow

[dan-zheng]

@swift-ci Please test tensorflow

[dan-zheng]

@swift-ci Please test tensorflow