[SR-9448] [AD] Diagnose @autodiff function types whose parameters or results don't conform to Differentiable

[rxwei]

Previous ID	SR-9448
Radar	None
Original Reporter	@rxwei
Type	Bug
Status	Resolved
Resolution	Done

Additional Detail from JIRA

Votes	0
Component/s	Swift for TensorFlow
Labels	Bug
Assignee	pschuh (JIRA)
Priority	Medium

md5: befdfabb4c9c4f81ed793d918da2b9d7

Issue Description:

Test case

struct NonDiffType { var x: Int }
// expected-error @+1 {{argument is not differentiable, but the enclosing function type is marked '@autodiff'; did you want to add '@nondiff' to this argument?}}
let _: @autodiff (NonDiffType) -> Float
// expected-error @+1 {{result is not differentiable, but the enclosing function type is marked '@autodiff'; did you want to add '@nondiff' to this argument?}}
let _: @autodiff (Float) -> NonDiffType

[rxwei]

pschuh (JIRA User) Interested in looking into this?

[swift-ci]

Comment by Parker Schuh (JIRA)

Hi Richard, I took a look at this and I'm able to generate the warnings in TypeCheckType.cpp, but it incorrectly fails on:

func evaldiff<T: Differentiable, U: Differentiable>(_ f: @autodiff (T) -> U, _ x: T) -> (U, (T.TangentVector) -> U.TangentVector)
where T == T.TangentVector {
return Builtin.autodiffApply_jvp(f, x)
}

I'm just trying to sort out how to check type conformance when it is a generic_type_param_type.

However, If I blindly allow all generics (as a hack) I an error for:

func methoddiff<T: Differentiable, U: Differentiable, R: Differentiable>(_ f: @autodiff (T) -> (U) -> R, _ x: T, _ y: U)

It is complaining about: (U) -> R not conforming to a protocol.

[rxwei]

Just sent #22166 for you to be able to calculate an AD-associated vector space of a function type.