Can't create a custom tangent vector type with computed properties (for lazy materialization)

[GeorgeLyon]

Description

Tangent vectors for aggregates (structs) can get really large, and there are many situations where you do not need to materialize all properties of the tangent vector in order to properly differentiate your code. Currently, there isn't a good way to do "lazy" materialization of a tangent vector because the compiler complains when a tangent vector's member is not a stored property (to be lazy, we'd need the property to be computed so that it can materialize itself on demand).

Here is an example that demonstrates the issue: https://github.com/GeorgeLyon/LazyTangentVectorIssue/blob/1f65e3333d0b6f303a60e43d358d2213e15925a0/Sources/DifferentiableSwiftExample/DifferentiableSwiftExample.swift

Reproduction

import _Differentiation

@differentiable(reverse)
public func foo(_ value: Foo) -> Double {
  return value.x + 3
}

public struct Foo: Differentiable {
  var x: Double

  public mutating func move(by offset: TangentVector) {
    self.x += offset.x
  }

  public struct TangentVector: Differentiable, AdditiveArithmetic {
    public var x: Double.TangentVector {
      get { _x }
      set { _x = newValue }
    }
    public var _x: Double.TangentVector
  }
}

Error:

error: function is not differentiable
 1 │ import _Differentiation
 2 │ 
 3 │ @differentiable(reverse)
   │  ╰─ error: function is not differentiable
 4 │ public func foo(_ value: Foo) -> Double {
   │             ╰─ note: when differentiating this function definition
 5 │   return value.x + 3
   │                ╰─ note: cannot differentiate access to property 'Foo.x' because 'Foo.TangentVector.x' is not a stored property
 6 │ }
 7 │ 

Expected behavior

This works.

Environment

Swift version 6.0-dev (LLVM ce41a43bba95b2b, Swift 1a84094)
Target: aarch64-unknown-linux-gnu
(swift-DEVELOPMENT-SNAPSHOT-2024-03-07-a)

Additional information

No response

[jkshtj]

Hey @dan-zheng, we have recently been trying to implement optimizations to lazily materialize a massive zero tangent vector for one of our types. One of the things that we tried, was to make the properties of the tangent vector “computed”, upon which the compiler gave us the following diagnostic.

note: cannot differentiate access to property 'Foo.x' because 'Foo.TangentVector.x' is not a stored property

It looks like you put those diagnostics in place. I wanted to understand whether this limitation is an artificial one. From briefly looking into it, it seems like we simply do not support this case because the pullback code generation becomes more complicated for us? Is that right or is there a deeper reason for the existence of this limitation?

[jkshtj]

@rxwei @asl would you guys have any insights about this?

[asl]

@GeorgeLyon Why closing?