[SR-13698] Conditionally conform Optional to Differentiable

[dan-zheng]

Previous ID	SR-13698
Radar	None
Original Reporter	@dan-zheng
Type	Sub-task
Status	Closed
Resolution	Done

Additional Detail from JIRA

Votes	1
Component/s
Labels	Sub-task
Assignee	@dan-zheng
Priority	Medium

md5: 13b8627652d2f04d11be34db0a4c54ae

Parent-Task:

• SR-13695 Enum differentiation

blocks:

• SR-13700 Differentiation transform: support optional-related operations

is duplicated by:

• TF-365 [AD] Make Optional conditionally conform to Differentiable

Issue Description:

Optional can conditionally conform to Differentiable when the Wrapped type does.

extension Optional : Differentiable where Wrapped : Differentaible {
  ...
}

Changes should be made in Optional.swift.

---

Prototype:

extension Optional : Differentiable where Wrapped : Differentiable {
  public enum DifferentiableView : Equatable, AdditiveArithmetic, Differentiable {
    case none
    case some(Wrapped.TangentVector)

    public static var zero: Self { .some(Wrapped.TangentVector.zero) }
    public static func + (lhs: Self, rhs: Self) -> Self {
      switch (lhs, rhs) {
        case (.none, .none): return .none
        case let (x, .none): return x
        case let (.none, y): return y
        case let (.some(x), .some(y)): return .some(x + y)
      }
    }

    public static func - (lhs: Self, rhs: Self) -> Self {
      switch (lhs, rhs) {
        case (.none, .none): return .none
        case let (x, .none): return x
        case let (.none, .some(y)): return .some(.zero - y)
        case let (.some(x), .some(y)): return .some(x - y)
      }
    }

    public typealias TangentVector = DifferentiableView
    public typealias AllDifferentiableVariables = Self
    public var allDifferentiableVariables: AllDifferentiableVariables {
      get { self }
      set { self = newValue }
    }
    public mutating func move(along direction: TangentVector) {
      switch (self, direction) {
      case (_, .none): return
      case let (.none, y): self = y
      case let (.some(x), .some(y)):
        var wrapped = x
        wrapped.move(along: y)
        self = .some(wrapped)
      }
    }
  }

  public typealias TangentVector = DifferentiableView
  public typealias AllDifferentiableVariables = Self
  public var allDifferentiableVariables: AllDifferentiableVariables {
    get { self }
    set { self = newValue }
  }
  public mutating func move(along direction: TangentVector) {
    switch (self, direction) {
    case (_, .none): return
    case (.none, _): fatalError("Move to move `.none`?")
    case let (.some(x), .some(y)):
      var wrapped = x
      wrapped.move(along: y)
      self = .some(wrapped)
    }
  }
}

Optional has a separate TangentVector type called DifferentiableView type, similar to Array.DifferentiableView. This is important to avoid conforming Optional to AdditiveArithmetic.

---

Note that differentiation of active Optional values is blocked by TF-583 (support for active enum values).

func loop_array(_ array: [Float]) -> Float {
  var result: Float = 1
  // for-in loop generates an active `Float?` value.
  for x in array {
    result = result * x
  }
  return result
}

[swift-ci]

Comment by Molly Beavers (JIRA)

Started working on this.

[dan-zheng]

Updated the issue with a prototype - sorry for not updating earlier.

[rxwei]

If the zero is represented as .some(.zero), we should fatalError on arithmetic cases that involve a nil and a non-nil.

[swift-ci]

Comment by Molly Beavers (JIRA)

I am still trying to figure this one out, though it has been two weeks so I thought I should give an update by at least this point.

I added the prototype code to Optional.swift and went to add tests in order to see what might need to be changed. I sort of copied how arrays(9c349704123fa4f2a783cafef0251e7299a1273b) and enums were tested below

OptionalTests.test("Optional Dense") {
 struct optDense : Differentiable {
   var w1: Float?

   @differentiable
   func callAsFunction(_ x: Float) -> Float {
     if let w1 = w1 {
      return x * w1
     }
     return x
   }
 }

 let dense=optDense(w1: 5)
 let grad = gradient(at: 3, in: { x in dense(x) })
 expectEqual(5.0,grad) //expected failure
}

The above fails with "differentiating enum values is not yet supported"

I then tried to follow the way enums were tested in TF-582

 OptionalTests.test("Optional Dense") {
  struct optDense : Differentiable {
   var w1: Float?
 
   @differentiable
   func callAsFunction(_ x: Float) -> Float {
     if let w1 = w1 {
       return x * w1
     }
     return x
   }
 }

 let grad = optDense(w1: 4).gradient(at: 2, in: { dense, x in dense(x) })
 expectEqual(0,1)
}

This fails with:
swift: /home/sftf/swift-source/swift/lib/SILGen/SILGenPoly.cpp:1444: void (anonymous namespace)::TranslateArguments::translateSingle(swift::Lowering::AbstractionPattern, swift::CanType, swift::Lowering::AbstractionPattern, swift::CanType, swift::Lowering::ManagedValue, s
wift::SILParameterInfo): Assertion `Outputs.back().getType() == SGF.getSILType(result)' failed.
I think this is because Optional is using a differential view?

I have also rebuilt swift/tensorflow quiet a few times now from scratch, so I do not believe it is a version issue. I have tried a bit more than this, but I always seem to come back to these two issues.

I have since been looking at SilGen and my next plan is to ask on the fastai forum to see if anyone else has run into similar issues. If I am making an obvious mistake is there a resource I can go to learn more in this area?

Currently using most recent code from swift/tensorflow. Can upload docker image if needed.

[rxwei]

Differentiating through enum switches is not possible yet because differentiating data-dependent control flow is not yet supported. So even if there's a fully correct Differentiable conformance, there is no way to test it against differentiation meaningfully. We'll update this issue once that is supported.

The crasher you encountered is a bug in the compiler. The DifferentiableView definition is not ideal: It should wrap an {{Optional}} instead.

[swift-ci]

Comment by Molly Beavers (JIRA)

Thank you for the help! I was stuck because I had a misunderstanding on the amount I should be expecting to change the prototype. Knowing that I can make significant changes allowed me to try a lot more.

"The DifferentiableView definition is not ideal: It should wrap an Optional instead." - Question: Did you mean the code from the prototype(Code #1) ->

public enum DifferentiableView : Equatable, AdditiveArithmetic, Differentiable {
 case none
 case some(Wrapped.TangentVector)

should instead be(Code #2)?:

public struct DifferentiableView : Differentiable {
 private var _base: Optional<Wrapped.TangentVector> //very similar pattern to Array.DifferentialView

I also tried(Code #3), Though I found that this did not work as Wrapped does not have to conform to Additive arithmetic, and therefore Wrapped is not guaranteed to have a valid zero.

public struct DifferentiableView : Differentiable {
 private var _base: Optional<Wrapped> //exactly like Array.DifferentialView

My current .zero code. (Code #4)

public static var zero: Optional<Wrapped>.DifferentiableView {
 let zero:Wrapped.TangentVector? = Wrapped.TangentVector.zero //Wrapped.zero not guaranteed
 return Optional.DifferentiableView(zero)
 )

I am currently working with the approach in (Code #2) and having results as I am able to test Differential view against AdditiveArithmetic tests before turning on Differentiation.

After going through and switching this to a struct I have found that "Assertion `Outputs.back().getType() == SGF.getSILType(result)' failed." still happens. Therefore I think you may have meant that (Code #1) should have instead been:

public enum DifferentiableView : Equatable, AdditiveArithmetic, Differentiable { 
case some(Wrapped)

So I am going to try and convert everything tomorrow.

[swift-ci]

Comment by Molly Beavers (JIRA)

Tried converting everything today, but not much luck as of yet. Unsure on how to debug compiler issues, is there any topics I can lookup to get started on learning?

Still running into the same "Assertion `Outputs.back().getType() == SGF.getSILType(result)' failed." issue. Continuing to attempt to find a way around it.

Keeeping code here: https://github.com/marii-moe/swift/tree/opt-scratch

[swift-ci]

Comment by Molly Beavers (JIRA)

Tried having the DifferentialView wrap a Optional<Wrapped>, and ran into the same compiler bug.

Currently attempting trying to debug the compiler referencing this: https://github.com/apple/swift/blob/master/docs/DebuggingTheCompiler.rst#debugging-the-swift-compiler

Not really familiar with C++, but I can hopefully at least be able to figure out more about the problem.

[swift-ci]

Comment by Molly Beavers (JIRA)

Upon checking SILGen I found that the issue seems to be that they are comparing the value to a pointer.

$*optDense.TangentVector
$optDense.TangentVector

Going to continue looking at this...

[swift-ci]

Comment by Molly Beavers (JIRA)

Still having issues getting debugging working. Seems SIGTRAP is being caught, so breakpoints are not working. Trying to figure out what to do here.

[swift-ci]

Comment by Molly Beavers (JIRA)

After looking through this a bit more I wanted to double check on something.

https://github.com/apple/swift/blob/tensorflow/stdlib/public/core/Array.swift

extension Array.DifferentiableView : AdditiveArithmetic
where Element : AdditiveArithmetic {
  public static func +
  public static func -
}

I have been struggling trying to determine if there was a way to do it with only Wrapped.TangentVector conforming to additive Arithmetic, to allow things like Optional<Optional<Wrapped>>. I figured this is why Wrapped did not conform to AdditiveArithmetic in the prototype. I am going to now work with the assumption that Wrapped conforming to AdditiveArithmetic is okay.

So, is it okay if Wrapped conforms to additive Arithmetic?

[swift-ci]

Comment by Molly Beavers (JIRA)

I now have something that looks like it is mostly working. Currently requesting access to the presentation today, because it went over what an "active Differentiable" means, which if I remember correctly looks like might be the issue with my tests.

[dan-zheng]

Hi Molly (JIRA User),

I'm so sorry for the late reply - thanks for investigating this issue!

Support for differentiating enums isn't quite robust (TF-583) and differentiation of enums hasn't been tested (it's not expected to work yet). I saw that you encountered some of these issues (e.g. Self vs TangentVector type mismatch, which is a differentiation transform bug/deficiency). Optional conformance to Differentiable isn't ready for robust testing at all, but proactively designing the conformance now (getting things to type-check) is valuable.

Please see the "Differentiable conformances" section of the Swift Differentiable Programming Design Overview for our latest ideas on what Differentiable conformances should look like for standard library types. A productive form of communication is the Swift for TensorFlow mailing list, if you'd like to start a thread on this issue - we'll try to be more responsive there. 🙂

[swift-ci]

Comment by Molly Beavers (JIRA)

@dan-zheng Thank you!

I am going to start a thread and update my code according to that document. I had previously been updating my code by referencing other code in the repository. I has skimmed the document previously, but did not realize the importance of it.

I can always update the code whenever active enums are supported.

[dan-zheng]

Thank you Molly (JIRA User), I look forward to your thread 🙂
I'm not sure when we'll visit enum differentiation - maybe within a month. I saw you're watching TF-583 - we'll post updates there.