[AutoDiff] Support inout argument differentiation.

[dan-zheng]

Includes cherry-pick of #29959: typing rules for inout parameter differentiation.

---

Add reverse-mode differentiation support for apply with inout arguments.

Notable pullback generation changes:

• If the pullback seed argument is inout, assign it (rather than a copy)
  directly as the adjoint buffer of the original result. This is important so
  the value is updated in-place.
• In visitApplyInst: skip adjoint accumulation for inout arguments.
  Adjoint accumulation for inout arguments occurs when callee pullbacks are
  applied, so no extra accumulation is necessary.

Add derivatives for functions with inout parameters in the stdlib for testing:

• FloatingPoint operations: +=, -=, *=, /=
• Array.append

Resolves TF-1165.

Todos:

• Add more tests, e.g. SILGen tests for inout derivative typing rules.
• Evaluate performance of inout derivatives vs functional derivatives + mutation.
• TF-1166: enable @differentiable attribute on set accessors.
• TF-1173: add forward-mode differentiation support for apply with inout
  parameters.

Exposes TF-1175: incorrect activity for class arguments.
Exposes TF-1176: incorrect activity for class modify accessors.
Add negative tests.

---

Examples:

// Test `inout` identity function.
func inoutIdentity(_ x: inout Float) {}
func identity(_ x: Float) -> Float {
  var result = x
  inoutIdentity(&result) // here
  return result
}
print(valueWithGradient(at: 10, in: identity))
// (value: 10.0, gradient: 1.0)

// Test `Float.*=` and control flow differentiation.
func product(_ array: [Float]) -> Float {
  var result: Float = 1
  for i in withoutDerivative(at: array.indices) {
    result *= array[i] // here
    // Old functional version:
    // result = result * array[i]
  }
  return result
}
print(valueWithGradient(at: [3, 4, 5], in: product))
// (value: 60.0, gradient: [20.0, 15.0, 12.0])

// Test `Array.append` and control flow differentiation.
func squared(_ array: [Float]) -> [Float] {
  var results: [Float] = []
  for i in withoutDerivative(at: array.indices) {
    results.append(array[i] * array[i]) // here
    // Old functional version:
    // results = results + [array[i] * array[i]]
  }
  return results
}
let v: Array<Float>.TangentVector = [4, -5, 6]
let (value, pb) = valueWithPullback(at: [1, 2, 3], in: squared)
print(value, pb(v))
// [1.0, 4.0, 9.0] [8.0, -20.0, 36.0]

// Test `set` accessor differentiation.
struct S: Differentiable {
  var x: Float

  var computed: Float {
    get { x }
    set { x = newValue }
  }
}
func squared(_ x: Float) -> Float {
  var s = S(x: 1)
  s.x *= x // here
  s.computed *= x // here
  return s.x
}
print(valueWithGradient(at: 10, in: squared))
// (value: 100.0, gradient: 20.0)

---

Thanks @marcrasi for peer debugging incorrect inout argument derivative values!

We discovered that Pullback::visitApplyInst should skip adjoint accumulation for inout arguments, which is interesting and not obvious. This made the first exciting test cases pass.

Working together was fun and productive 🙂

[dan-zheng]

Note: merging master branch PR #29959 (inout differentiation typing rules) to tensorflow branch requires this PR, which has non-trivial accompanying changes to fix tests:

Failing Tests (5):
    Swift(linux-x86_64) :: AutoDiff/downstream/activity_analysis.swift
    Swift(linux-x86_64) :: AutoDiff/downstream/differentiation_transform_diagnostics.swift
    Swift(linux-x86_64) :: AutoDiff/downstream/differentiable_attr_type_checking.swift
    Swift(linux-x86_64) :: AutoDiff/downstream/forward_mode_diagnostics.swift
    Swift(linux-x86_64) :: AutoDiff/compiler_crashers_fixed/tf1167-differentiable-attr-override-match.swift

This is blocking the master -> tensorflow merge, so I'd like to land this PR as soon as possible.

Lesson learned: next time when doing changes that touch both upstream and downstream code, create PRs to both tensorflow and master branch, and test+merge the tensorflow branch PR first. This makes the master -> tensorflow merge process much easier.

[saeta]

🎉

[rxwei]

Missing test cases:

1. Differentiating a function with inout parameters that have a class type. This is expected to work normally.
2. Differentiating a function whose body calls a same-file function that both takes a Differentiable inout argument and returns a Differentiable formal result where both the inout argument and the formal result are active. This is expected to trigger a differentiation error.

[rxwei]

Any forward mode tests?

[dan-zheng]

Any forward mode tests?

Forward-mode differentiation doesn't support inout arguments yet: TF-1173.
Adding support should involve straightforward changes to JVPEmitter::visitApplyInst.
In the meantime, forward-mode inout argument differentiation remains diagnosed.

[dan-zheng]

@swift-ci Please test tensorflow