Skip to content

[AutoDiff] Initial Forward Mode AD Support #26057

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 51 commits into from
Aug 20, 2019

Conversation

bartchr808
Copy link
Contributor

@bartchr808 bartchr808 commented Jul 10, 2019

This PR extends on the previous work here on the JVPEmitter (#25954) to make it be able to generate SIL code for the differential function, as well as more correctly implementing the body of the JVP.'

generates a valid JVP and differential. I have yet to add support for control flow and generics. This PR also renames PullbackInfo to LinearMap info since it has a lot of tangent mapping helpers that JVPEmitter can use when creating the differential.

This is still early, so the JVP/Differential emission only runs with the flag -Xllvm -run-jvp-generation when running Swift. With it, you can differentiate basic Float functions with + - / * with no vars like:

@differentiable
func foo(x: Float, y: Float) -> Float {
  return x * x * x * y
}

@differentiable
func binary(x: Float, y: Float) -> Float {
  let a = y / x
  return x * y + a
}

@bartchr808 bartchr808 added the tensorflow This is for "tensorflow" branch PRs. label Jul 10, 2019
@bartchr808 bartchr808 force-pushed the differential-emitter branch from 04be590 to 0ca031b Compare July 11, 2019 10:35
@bartchr808 bartchr808 force-pushed the differential-emitter branch from 40b3680 to 652a6d7 Compare July 11, 2019 13:44
@bartchr808 bartchr808 force-pushed the differential-emitter branch from 9835cb0 to 966c470 Compare July 12, 2019 16:01
@bartchr808 bartchr808 force-pushed the differential-emitter branch from 3c01ab8 to 5883668 Compare July 15, 2019 20:58
@bartchr808 bartchr808 marked this pull request as ready for review July 16, 2019 22:03
@bartchr808 bartchr808 requested review from rxwei and dan-zheng July 16, 2019 22:14
Copy link
Contributor

@dan-zheng dan-zheng left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you please resolve merge conflicts and rebase into a single commit?
EDIT: "Squash and merge" instead of rebasing is fine as well.

@bartchr808 bartchr808 requested review from rxwei and dan-zheng August 12, 2019 18:34
@bartchr808
Copy link
Contributor Author

@swift-ci please test tensorflow

1 similar comment
@bartchr808
Copy link
Contributor Author

@swift-ci please test tensorflow

@bartchr808 bartchr808 changed the title [AutoDiff] Differential emitter [AutoDiff] Initial Forward Mode AD Support Aug 12, 2019
@bartchr808 bartchr808 force-pushed the differential-emitter branch from 051a924 to 4e60732 Compare August 19, 2019 23:54
@bartchr808 bartchr808 force-pushed the differential-emitter branch from 4e60732 to 1883ac1 Compare August 19, 2019 23:55
@bartchr808
Copy link
Contributor Author

For time going to run tests now to get this merged in as soon as possible assuming there are only NFC changes that need to be made

@bartchr808
Copy link
Contributor Author

@swift-ci please test tensorflow

1 similar comment
@bartchr808
Copy link
Contributor Author

@swift-ci please test tensorflow

@bartchr808 bartchr808 force-pushed the differential-emitter branch from 6007902 to 129a3d5 Compare August 20, 2019 03:34
@bartchr808
Copy link
Contributor Author

@swift-ci please test tensorflow

1 similar comment
@bartchr808
Copy link
Contributor Author

@swift-ci please test tensorflow


ForwardModeTests.test("Protocols") {
let inst = Linear(m: 5, b: -2)
let (y1, diff1) = valueWithDifferential(at: 5) { x in inst.foo(x: x) }
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is not testing differentiation through protocols at all because you calling a method with concrete Self type Linear.

To differentiate through a protocol requirement, define a generic function that takes an instance bound to a generic parameter that conforms to the protocol, and differentiate that.

func genericFoo<T: DiffReq>(_ t: T, _ x: Float) -> Float {
    t.foo(x)
}
...
valueWithDifferential(at: 5) { x in genericFoo(inst, x) }

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

discussed offline, will address this in the next PR which will support everything else

@bartchr808
Copy link
Contributor Author

@swift-ci please test tensorflow

2 similar comments
@bartchr808
Copy link
Contributor Author

@swift-ci please test tensorflow

@bartchr808
Copy link
Contributor Author

@swift-ci please test tensorflow

Copy link
Contributor

@rxwei rxwei left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Congrats on going through all the hurdles to finish this important step!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
tensorflow This is for "tensorflow" branch PRs.
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants