// file1.swift
import _Differentiation

protocol P1: Differentiable {
  @differentiable(wrt: self)
  func callAsFunction(_ input: Float) -> Float
}

protocol P2: P1 {}

extension P2 {
  // Adding this `@differentiable` attribute fixes the issue:
  // @differentiable(wrt: self)
  @differentiable(wrt: (self, input))
  public func callAsFunction(_ input: Float) -> Float {
    return input
  }
}

// file2.swift
import _Differentiation

struct S: P2 {}

$ swiftc -emit-library file1.swift file2.swift -module-name main
Undefined symbols for architecture x86_64:
  "_AD__$s4main2P2PAAE14callAsFunctionyS2fF_PUSRS4main2P2Rzl", referenced from:
      _AD__$s4main1SVAA2P1A2aDP14callAsFunctionyS2fFTW_jvp_US in file2-c8cbd9.o
      _AD__$s4main1SVAA2P1A2aDP14callAsFunctionyS2fFTW_vjp_US in file2-c8cbd9.o
ld: symbol(s) not found for architecture x86_64
<unknown>:0: error: link command failed with exit code 1 (use -v to see invocation)

// AD__$s4main1SVAA2P1A2aDP14callAsFunctionyS2fFTW_jvp_US
sil private [transparent] [thunk] @AD__$s4main1SVAA2P1A2aDP14callAsFunctionyS2fFTW_jvp_US : $@convention(witness_method:
 P1) (Float, @in_guaranteed S) -> (Float, @owned @callee_guaranteed @substituted <τ_0_0> (@in_guaranteed τ_0_0) -> Float
 for <S.TangentVector>) {
// %0                                             // user: %7
// %1                                             // user: %7
bb0(%0 : $Float, %1 : $*S):
  // function_ref P2.callAsFunction(_:)
  %2 = function_ref @$s4main2P2PAAE14callAsFunctionyS2fF : $@convention(method) <τ_0_0 where τ_0_0 : P2> (Float, @in_gua
ranteed τ_0_0) -> Float // user: %5
  %3 = differentiability_witness_function [jvp] [parameters 1] [results 0] <Self where Self : P2> @$s4main2P2PAAE14callA
sFunctionyS2fF : $@convention(method) <τ_0_0 where τ_0_0 : P2> (Float, @in_guaranteed τ_0_0) -> Float // user: %5
  %4 = differentiability_witness_function [vjp] [parameters 1] [results 0] <Self where Self : P2> @$s4main2P2PAAE14callA
sFunctionyS2fF : $@convention(method) <τ_0_0 where τ_0_0 : P2> (Float, @in_guaranteed τ_0_0) -> Float // user: %5
  %5 = differentiable_function [parameters 1] [results 0] %2 : $@convention(method) <τ_0_0 where τ_0_0 : P2> (Float, @in
_guaranteed τ_0_0) -> Float with_derivative {%3 : $@convention(method) <τ_0_0 where τ_0_0 : P2> (Float, @in_guaranteed τ
_0_0) -> (Float, @owned @callee_guaranteed @substituted <τ_0_0> (@in_guaranteed τ_0_0) -> Float for <τ_0_0.TangentVector
>), %4 : $@convention(method) <τ_0_0 where τ_0_0 : P2> (Float, @in_guaranteed τ_0_0) -> (Float, @owned @callee_guarantee
d @substituted <τ_0_0> (Float) -> @out τ_0_0 for <τ_0_0.TangentVector>)} // user: %6
  %6 = differentiable_function_extract [jvp] %5 : $@differentiable @convention(method) <τ_0_0 where τ_0_0 : P2> (@noDeri
vative Float, @in_guaranteed τ_0_0) -> Float // user: %7
  %7 = apply %6<S>(%0, %1) : $@convention(method) <τ_0_0 where τ_0_0 : P2> (Float, @in_guaranteed τ_0_0) -> (Float, @own
ed @callee_guaranteed @substituted <τ_0_0> (@in_guaranteed τ_0_0) -> Float for <τ_0_0.TangentVector>) // user: %8
  return %7 : $(Float, @callee_guaranteed @substituted <τ_0_0> (@in_guaranteed τ_0_0) -> Float for <S.TangentVector>) //
 id: %8
} // end sil function 'AD__$s4main1SVAA2P1A2aDP14callAsFunctionyS2fFTW_jvp_US'

// file1.swift
import _Differentiation

protocol P1: Differentiable {
  @differentiable(wrt: self)
  func callAsFunction(_ input: Float) -> Float
}

protocol P2: P1 {
  @differentiable(wrt: (self, input))
  func callAsFunction(_ input: Float) -> Float

  @differentiable(wrt: (self, input))
  func forward(_ input: Float) -> Float
}

extension P2 {
  @differentiable(wrt: (self, input))
  public func forward(_ input: Float) -> Float {
    return callAsFunction(input)
  }

  @differentiable(wrt: (self, input))
  public func callAsFunction(_ input: Float) -> Float {
    return forward(input)
  }
}

// file2.swift
import _Differentiation

public struct S: P2 {
  @differentiable
  func forward(_ input: Float) -> Float {
    input
  }
}