[AutoDiff] Relax @differentiable requirement for protocol witnesses.

[dan-zheng]

Previously, all witnesses of a @differentiable protocol requirement were
required to have the same attribute (or one with superset parameter indices).

However, this leads to many annotations on witnesses and is not ideal for
usability. @differentiable attributes are really only significant on
public witnesses, so that they are clearly @differentiable at a glance (in
source code, interface files, and API documentation), without looking through
protocol conformance hierarchies.

Now, only public witnesses of @differentiable protocol requirements are
required to have the same attribute (or one with superset parameter indices).
For less-visible witnesses, an implicit @differentiable attribute is created
with the same configuration as the requirement's.

Resolves TF-1117.

---

This usability improvement was discussed during the
2020-01-17 Swift for TensorFlow open design review.
We agreed that it's a good idea!

Todo: upstream changes to master.

---

Example:

public protocol Layer: Differentiable {
  associatedtype Input: Differentiable
  associatedtype Output: Differentiable
  @differentiable(wrt: (self, input))
  func callAsFunction(_ input: Input) -> Output
}

// Internal conforming type.
struct DummyInternalLayer: Layer {
  func callAsFunction(_ input: Float) -> Float {
    return input
  }
}

Before (misleading diagnostic due to TF-1014):

layer.swift:9:8: error: type 'DummyInternalLayer' does not conform to protocol 'Layer'
struct DummyInternalLayer: Layer {
       ^
layer.swift:2:18: note: protocol requires nested type 'Input'; do you want to add it?
  associatedtype Input: Differentiable
                 ^
layer.swift:3:18: note: protocol requires nested type 'Output'; do you want to add it?
  associatedtype Output: Differentiable
                 ^

After: no error.

An implicit @differentiable attribute is created for DummyInternalLayer.callAsFunction:

$ swiftc -print-ast layer.swift
...

internal struct DummyInternalLayer : Layer {
  @differentiable(wrt: (self, input))
  internal func callAsFunction(_ input: Float) -> Float
  ...
}

[rxwei]

This makes the implementation consistent with #29307. LGTM.

[dan-zheng]

@swift-ci Please test tensorflow

[rxwei]

One thing: Creating implicit attributes leads to invalid source locations. Have you handled such @differentiable attributes differentiation diagnostics? In the transform, the differentiation invoker will be DifferentiableAttribute, and any diagnostics will try to get the source location of the implicit attribute.

[marcrasi]

I thought of a possible weird edge case. What if the conformance is declared in a separate file from the file where the function satisfying the requirement is declared? Then the differentiation pass on the file with the function never sees the implicit differentiable attribute because it only exists in the separate compilation where the conformance is declared.

This edge case probably also applies to the existing subset logic, so it's probably an existing problem that can be dealt with later.

[rxwei]

@marcrasi Interesting case. One possible solution is to disallow any implicit attribute inheritance when the witness is in a separate translation unit. We’ll just emit an error in protocol witness matching, stating which attribute the potential witness is missing.

[dan-zheng]

I thought of a possible weird edge case. What if the conformance is declared in a separate file from the file where the function satisfying the requirement is declared? Then the differentiation pass on the file with the function never sees the implicit differentiable attribute because it only exists in the separate compilation where the conformance is declared.

I added cross-file tests for this edge case:

• main.swift
• other_file.swift

Indeed, implicit @differentiable attributes cannot be generated for protocol witnesses when the conformance is declared from a separate file from the witnesses. Otherwise, compilation of the file containing the conformance creates external references to symbols for implicit @differentiable attributes, even though no such symbols exist.

I implemented @rxwei's solution: a diagnostic is emitted for these cases.

---

One thing: Creating implicit attributes leads to invalid source locations. Have you handled such @differentiable attributes differentiation diagnostics? In the transform, the differentiation invoker will be DifferentiableAttribute, and any diagnostics will try to get the source location of the implicit attribute.

Implicit protocol witness @differentiable attributes are actually created with the source location from the protocol requirement's @differentiable attribute.

I added a test to verify this:

protocol Protocol: Differentiable {
  @differentiable(wrt: (self, x))
  func internalMethod(_ x: Float) -> Float
}

struct ConformingStruct: Protocol {
  // Non-differentiable body!
  func internalMethod(_ x: Float) -> Float {
    return Float(Int(x))
  }
}

test.swift:2:4: error: function is not differentiable
  @differentiable(wrt: (self, x))
  ~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
test.swift:8:8: note: when differentiating this function definition
  func internalMethod(_ x: Float) -> Float {
       ^
test.swift:9:18: note: cannot differentiate through a non-differentiable result; do you want to use 'withoutDerivative(at:)'?
    return Float(Int(x))
                 ^

I think these source locations are reasonable. Let me know if you have better suggestions.

---

Ready for re-review! Changes will be upstreamed to master.

[rxwei]

I thought of a possible weird edge case. What if the conformance is declared in a separate file from the file where the function satisfying the requirement is declared? Then the differentiation pass on the file with the function never sees the implicit differentiable attribute because it only exists in the separate compilation where the conformance is declared.

I added cross-file tests for this edge case:

• main.swift
• other_file.swift

Indeed, implicit @differentiable attributes cannot be generated for protocol witnesses when the conformance is declared from a separate file from the witnesses. Otherwise, compilation of the file containing the conformance creates external references to symbols for implicit @differentiable attributes, even though no such symbols exist.

I implemented @rxwei's solution: a diagnostic is emitted for these cases.

One thing: Creating implicit attributes leads to invalid source locations. Have you handled such @differentiable attributes differentiation diagnostics? In the transform, the differentiation invoker will be DifferentiableAttribute, and any diagnostics will try to get the source location of the implicit attribute.

Implicit protocol witness @differentiable attributes are actually created with the source location from the protocol requirement's @differentiable attribute.

I added a test to verify this:

protocol Protocol: Differentiable {
  @differentiable(wrt: (self, x))
  func internalMethod(_ x: Float) -> Float
}

struct ConformingStruct: Protocol {
  // Non-differentiable body!
  func internalMethod(_ x: Float) -> Float {
    return Float(Int(x))
  }
}

test.swift:2:4: error: function is not differentiable
  @differentiable(wrt: (self, x))
  ~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
test.swift:8:8: note: when differentiating this function definition
  func internalMethod(_ x: Float) -> Float {
       ^
test.swift:9:18: note: cannot differentiate through a non-differentiable result; do you want to use 'withoutDerivative(at:)'?
    return Float(Int(x))
                 ^

I think these source locations are reasonable. Let me know if you have better suggestions.

Ready for re-review! Changes will be upstreamed to master.

It is great that we don't have invalid source locations, but this error message could be improved -- it's currently showing an error on the protocol definition as if the protocol was defined wrong, which is misleading. Instead, we should show an error on the declaration that actually triggers the error, along with a note that points us to the protocol requirement that had the @differentiable attribute.

[dan-zheng]

It is great that we don't have invalid source locations, but this error message could be improved -- it's currently showing an error on the protocol definition as if the protocol was defined wrong, which is misleading. Instead, we should show an error on the declaration that actually triggers the error, along with a note that points us to the protocol requirement that had the @differentiable attribute.

Done in 03a22e5:

test.swift:20:8: error: function is not differentiable
  func internalMethod(_ x: Float) -> Float {
       ^~~~~~~~~~~~~~
test.swift:20:8: note: when differentiating this function definition
  func internalMethod(_ x: Float) -> Float {
       ^
test.swift:22:18: note: cannot differentiate through a non-differentiable result; do you want to use 'withoutDerivative(at:)'?
    return Float(Int(x))
                 ^

Ready for re-review.

[rxwei]

Hmm the updated implementation is not showing a note on the protocol requirement that had the attribute.

[dan-zheng]

Hmm the updated implementation is not showing a note on the protocol requirement that had the attribute.

Done:

test.swift:17:15: error: type 'PublicConformingStruct' does not conform to protocol 'InternalProtocol'
public struct PublicConformingStruct: InternalProtocol {
              ^
test.swift:20:15: note: candidate is missing attribute '@differentiable'
  public func publicMethod(_ x: Float) -> Float {
              ^
         @differentiable
test.swift:8:4: note: required '@differentiable' attribute declared here
  @differentiable(wrt: (self, x))
  ~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
test.swift:9:8: note: protocol requires function 'publicMethod' with type '(Float) -> Float'; do you want to add a stub?
  func publicMethod(_ x: Float) -> Float
       ^

The last diagnostic is misleading because function publicMethod has been declared, but changing/removing it seems non-trivial. I'd like to defer further diagnostic improvements: TF-1154.

[rxwei]

I was not referring to the attribute-missing diagnostic. I was referring to the differentiation error.

test.swift:20:8: error: function is not differentiable
  func internalMethod(_ x: Float) -> Float {
       ^~~~~~~~~~~~~~
test.swift:20:8: note: when differentiating this function definition
  func internalMethod(_ x: Float) -> Float {
       ^
test.swift:22:18: note: cannot differentiate through a non-differentiable result; do you want to use 'withoutDerivative(at:)'?
    return Float(Int(x))
                 ^

This error should show a note that points to the trigger of differentiation, aka. the @differentiable attribute on the protocol requirement that got implicitly inherited.

[dan-zheng]

This error should show a note that points to the trigger of differentiation, aka. the @differentiable attribute on the protocol requirement that got implicitly inherited.

Oops, sorry for the misunderstanding.

I reverted a commit in the wrong direction.
Implicitly inherited @differentiable attributes now have a dedicated diagnostic:

protocol Protocol: Differentiable {
  @differentiable(wrt: (self, x))
  func internalMethod(_ x: Float) -> Float
}

struct ConformingStruct: Protocol {
  func internalMethod(_ x: Float) -> Float {
    return Float(Int(x))
  }
}

test.swift:7:8: error: function is not differentiable
  func internalMethod(_ x: Float) -> Float {
       ^~~~~~~~~~~~~~
test.swift:2:4: note: differentiability required by the corresponding protocol requirement here
  @differentiable(wrt: (self, x))
   ^~~~~~~~~~~~~~
test.swift:8:18: note: cannot differentiate through a non-differentiable result; do you want to use 'withoutDerivative(at:)'?
    return Float(Int(x))
                 ^

[dan-zheng]

Note: it's possible to generalize this "diagnostic source location" for other kinds of implicit @differentiable attributes too, not just "ones that are implicitly inherited from protocol requirements".

Currently, there are only three kinds of implicit @differentiable attributes:

1. This case: @differentiable attributes inherited from protocol requirements.
2. @differentiable attributes synthesized from others on the same declaration with superset parameter indices.
   • @differentiable(wrt: (x, y)) -> @differentiable(wrt: x)
3. @differentiable attributes synthesized on stored properties of structs/classes that derive a conformance to Differentiable.

enum class ImplicitKind {
  InheritedFromProtocolRequirement, // this case
  SynthesizedSubsetParametersAttribute, // triggered during conformance type-checking
  SynthesizedForStoredProperty, // triggered during `Differentiable` derived conformances
};

I think the first case benefits the most from a "secondary diagnostic source location". Generalization could be done later.

[rxwei]

Thanks for addressing all the comments!

[dan-zheng]

Thanks for the thorough review!
I'll try upstreaming changes to master first, then merging the changes into tensorflow.

[dan-zheng]

@swift-ci Please test tensorflow

[dan-zheng]

I accidentally deleted tensorflow branch, which closed this PR. That was not intentional, sorry!
It would be nice to protect tensorflow branch against deletion to prevent this from happening again.

[dan-zheng]

@swift-ci Please test tensorflow

[dan-zheng]

@swift-ci Please test tensorflow