[AutoDiff] [stdlib] Made arrays differentiable #23183

eaplatanios · 2019-03-08T19:03:56Z

This is a first attempt at making arrays differentiable.

stdlib/public/core/Array.swift

…ons.

stdlib/public/core/Array.swift

eaplatanios · 2019-03-08T21:48:30Z

Thanks Richard! I addressed all the feedback and tests seem to pass locally on my machine! :)

rxwei · 2019-03-08T22:01:02Z

Have you checked in those tests? I don’t see them in this PR. You can create a new file: test/AutodDiff/array.swift for tests.

eaplatanios · 2019-03-08T22:09:30Z

Sorry I haven't. I'll add a few more and push in a few minutes. Quick question: is there a way to recompile and run only the newly added tests? utils/build-script --preset=tensorflow_test seems to be recompiling a lot of things and running all tests, which is killing my MacBook and is also taking forever every time I want to test.

marcrasi · 2019-03-08T22:17:04Z

Quick question: is there a way to recompile and run only the newly added tests?

Yes!

To do an incremental compile, cd to build/<build dir name>/swift-<something> and run ninja. (build should be inside the parent directory of the swift repo. <build dir name> depends on what settings you used to build, and there's probably only one, unless you have built with multiple different settings.)

You can also run ninja swift to just recompile the compiler without recompiling the stdlib, to save time in cases where your changes don't affect the stdlib. But this PR is a change to the stdlib.

To run a specific test, you need to construct a lit invocation that runs that test. There's some documentation about it here: https://github.com/apple/swift/blob/master/docs/Testing.md. But doing it manually is pretty tedious, so I created a script where you can just pass it the test file and it'll automatically run lit in the correct way: https://gist.github.com/marcrasi/7493150a7247b0a43579126d3a6352de (I should really check this script into a repo and document it somewhere...)

eaplatanios · 2019-03-08T22:31:33Z

Thanks a lot Marc! That's extremely useful! :) I made a couple simple tests, including the example you mentioned in the discussion group. I'll make sure they pass and commit them.

stdlib/public/core/Array.swift

rxwei · 2019-03-09T01:11:30Z

LGTM, thanks for pushing this through! If you can check in the tests, that'd be great!

eaplatanios · 2019-03-09T01:17:13Z

Thanks for all the help! I'm about to commit a couple simple tests, but I can't seem to run them using the litrun script by @marcrasi . Marc, I'm getting the following error: RuntimeError: Device tests are currently only supported when the swift_test_mode is "only_non_executable". Current swift_test_mode is optimize_none.. Do you have any experience with that? I'm trying to run them from the build directory by using: ./litrun -b Ninja-RelWithDebInfoAssert ../../../swift/test/AutoDiff/array.swift.

saeta · 2019-04-16T23:31:04Z

Thanks for the update @marcrasi . re:#2: That's exactly the semantics that I implemented for the LayerList (that's a workaround for not having differentiable arrays.)

marcrasi · 2019-04-17T04:12:37Z

Here's an updated implementation that takes @rxwei's comments (above + some in-person discussion) into account: https://github.com/fastai/fastai_docs/blob/master/dev_swift/01c_array_differentiable.ipynb . (Also viewable as gist if the github ipynb viewer is broken: https://gist.github.com/marcrasi/b7618e3e7a8a8a920b73e662a40c4c7b#file-array-differentiable-swift)

@eaplatanios, what do you think about this approach?

eaplatanios · 2019-04-17T15:56:26Z

@marcrasi Thanks lot for the update! I'll go through it later today, but just a quick question that pops in my mind is: what would the VJP of the .append function look like?

rxwei · 2019-04-17T17:13:18Z

append will require AD to support for mutating functions and inout arguments. Roughly, AD will be treating an inout argument (or a mutating self) as both an input and an output.

eaplatanios · 2019-04-17T17:32:36Z

@rxwei assuming that AD supports mutating functions and inout arguments, what would be the array length for the gradient?

marcrasi · 2019-04-17T17:51:33Z

We can think of concat as the immutable version of append, so a VJP for concat demonstrates how a VJP for append can eventually work:

// We should actually make the stdlib "+" differentiable instead of defining our own "concat". But it was
// easier to define my own "concat" for the example.
@differentiable(vjp: vjpConcat where E: Differentiable)
func concat<E>(_ lhs: [E], _ rhs: [E]) -> [E] {
    return lhs + rhs
}

func vjpConcat<E: Differentiable>(_ lhs: [E], _ rhs: [E])
    -> ([E], (Array<E>.CotangentVector) -> (Array<E>.CotangentVector, Array<E>.CotangentVector))
{
    func pullback(_ v: Array<E>.CotangentVector) -> (Array<E>.CotangentVector, Array<E>.CotangentVector) {
        precondition(
            v.base.count == lhs.count + rhs.count,
            "concat received gradient with wrong count. gradient: \(v.base.count), lhs: \(lhs.count), rhs: \(rhs.count)")
        // It would be more efficient if we find a way to keep the slices as ArraySlices.
        return (
            Array<E>.CotangentVector(Array(v.base[0..<lhs.count])),
            Array<E>.CotangentVector(Array(v.base[lhs.count...])))
    }
    return (concat(lhs, rhs), pullback)
}

// Similarly, we should actually make the stdlib subscript differentiable.
extension Array where Element : Differentiable {
  @differentiable(wrt: (self), vjp: vjpDifferentiableSubscript)
  func differentiableSubscript(_ i: Int) -> Element {
    return self[i]
  }

  func vjpDifferentiableSubscript(_ i: Int) -> (Element, (Element.CotangentVector) -> Array.CotangentVector) {
    let result = self[i]
    func pullback(_ v: Element.CotangentVector) -> Array.CotangentVector {
      var r = Array<Element.CotangentVector>(repeating: .zero, count: count)
      r[i] = v
      return Array.CotangentVector(r)
    }
    return (result, pullback)
  }
}

@differentiable
func sumFirstThreeConcatted(_ arrs: TwoArrs) -> Float {
    let c = concat(arrs.a, arrs.b)
    return c.differentiableSubscript(0) + c.differentiableSubscript(1) + c.differentiableSubscript(2)
}

struct TwoArrs: Differentiable {
    let a: [Float]
    let b: [Float]
    init(_ a: [Float], _ b: [Float]) {
        self.a = a
        self.b = b
    }
}

print(gradient(at: TwoArrs([0, 0], [0, 0]), in: sumFirstThreeConcatted))
// => (a: [1.0, 1.0], b: [1.0, 0.0])

print(gradient(at: TwoArrs([], [0, 0, 0]), in: sumFirstThreeConcatted))
// => (a: [], b: [1.0, 1.0, 1.0])

print(gradient(at: TwoArrs([0, 0, 0], []), in: sumFirstThreeConcatted))
// => (a: [1.0, 1.0, 1.0], b: [])

The append VJP should do something similar. It's pullback will take a gradient of length n + 1, "return" the first n elements as the gradient of the input array, and return the last element as the gradient of the appended element.

marcrasi · 2019-04-17T18:00:34Z

@eaplatanios, do you mind if I update this PR with these changes and try to get it merged soon? We're cutting a "v0.3" release of S4TF this evening, and we'd really like to have differentiable arrays in it!

…/swift-language into eaplatanios-array-differentiable

eaplatanios · 2019-04-17T18:47:04Z

@marcrasi Sorry it's been a super busy day so far. I took a look at your update and I agree it makes sense and like how we also avoid a lot of code duplication. Plus it makes more sense mathematically (I really like the last example you added in your earlier post). Feel free to update the PR and try to get it merged soon. I'm excited to see this making it in the 0.3 release. :)

rxwei · 2019-04-17T20:44:25Z

Let me know when it's ready to review!

marcrasi · 2019-04-17T20:44:56Z

@rxwei This is ready to review!

marcrasi · 2019-04-17T20:53:00Z

Note about thing I ran into: I also got the crash that @eaplatanios had. It seems related to cross-module stuff, which is why I didn't get the crash when I was playing with the gists. I fixed it by adding @_fixed_layout to DifferentiableView. Seems like there is an underlying bug, so I filed one here: https://bugs.swift.org/browse/TF-442 . This bug doesn't block this PR because @_fixed_layout makes it work.

rxwei · 2019-04-17T20:45:58Z

stdlib/public/core/Array.swift

+    // TODO: Determine if that is a bug, and fix.
+    public var array: [Element] {
+      @differentiable(wrt: self, vjp: _vjpArray)
+      get {


Use _read and _modify accessors instead of get/set to make it more efficient.

Suggested change

get {

_read { yield base }

_modify { yield &base }

The @differentiable attr doses not like _read. It says:

/usr/local/google/home/marcrasi/swift-base/swift/stdlib/public/core/Array.swift:1888:8: error: cannot differentiate void function '_' @differentiable(wrt: self, vjp: _vjpBase) ^

So I will leave it as get / _modify for now, and this whole problem will go away once we fix the @differentiable stored property issue.

rxwei · 2019-04-17T20:48:41Z

stdlib/public/core/Array.swift

+    // I'm implementing this as a computed property instead of directly exposing `base` because the
+    // `@differentiable` annotation does not make the stored property actually differentiable. I
+    // think this is a bug. Maybe it's related to `@_fixed_layout`?
+    // TODO: Determine if that is a bug, and fix.


Interesting. How about renaming the current base to _base, and renaming array/_vjpArray to base/_vjpBase?

rxwei · 2019-04-17T20:49:35Z

stdlib/public/core/Array.swift

+
+    /// Construct a view of the given array.
+    @differentiable(wrt: array, vjp: _vjpInit)
+    public init(_ array: [Element]) { self.base = array }


Nit:

Suggested change

public init(_ array: [Element]) { self.base = array }

public init(_ base: [Element]) { self.base = base }

rxwei · 2019-04-17T20:49:44Z

stdlib/public/core/Array.swift

+      return (array, { $0 })
+    }
+
+    /// Construct a view of the given array.


Suggested change

/// Construct a view of the given array.

/// Creates a differentiable view of the given array.

rxwei · 2019-04-17T20:50:22Z

stdlib/public/core/Array.swift

+
+// SWIFT_ENABLE_TENSORFLOW
+extension Array where Element : Differentiable {
+  /// Views the array as the differentiable product manifold of `Element` multiplied with itself


Suggested change

/// Views the array as the differentiable product manifold of `Element` multiplied with itself

/// The view of an array as the differentiable product manifold of `Element` multiplied with itself

rxwei · 2019-04-17T20:52:24Z

stdlib/public/core/Array.swift

+      get {
+        return AllDifferentiableVariables(array.map { $0.allDifferentiableVariables })
+      }
+      set(newValue) {


Omit newValue since the name is implicitly bound.

Suggested change

set(newValue) {

set {

rxwei · 2019-04-17T20:53:40Z

stdlib/public/core/Array.swift

+    public func tangentVector(from cotangentVector: CotangentVector) -> TangentVector {
+      precondition(
+        array.count == cotangentVector.array.count,
+        "cannot use Array.DifferentiableView with count \(array.count) to get tangentVector from " +


Are these over 80 columns?

oops, I thought the limit was 100 based on some existing lines that go over 80

Yes sorry that was my mistake. I assumed it was 100 because that's in the Google Swift Style Guide.

rxwei · 2019-04-17T20:55:25Z

stdlib/public/core/Array.swift

+    get {
+      return DifferentiableView(self).allDifferentiableVariables
+    }
+    set(v) {


Suggested change

set(v) {

set {

Use newValue?

rxwei · 2019-04-17T20:57:43Z

stdlib/public/core/KeyPathIterable.swift

@@ -104,6 +104,13 @@ extension Array : KeyPathIterable {
  }
 }

+extension Array.DifferentiableView : KeyPathIterable {


I don't think you need a special implementation here. Derived conformances should just do it for you when you declare : KeyPathIterable at DifferentiableView's declaration site.

rxwei · 2019-04-17T20:58:55Z

stdlib/public/core/Array.swift

+      }
+    }
+
+    public func _vjpArray() -> ([Element], (Array<Element>.CotangentVector) -> CotangentVector) {


Make all VJPs be @usableFromInline internal

rxwei · 2019-04-17T21:00:39Z

The crash could be caused by @differentiable function lowering / VJP type calculation that we have today.

marcrasi · 2019-04-17T22:29:53Z

@swift-ci please test tensorflow

marcrasi · 2019-04-17T22:30:15Z

@swift-ci please test tensorflow

marcrasi · 2019-04-17T22:30:20Z

@swift-ci please test tensorflow

marcrasi · 2019-04-17T22:30:43Z

@swift-ci please test tensorflow

marcrasi · 2019-04-17T22:31:06Z

@swift-ci please test tensorflow

marcrasi · 2019-04-17T22:41:49Z

@swift-ci please test tensorflow

rxwei · 2019-04-23T05:46:45Z

@eaplatanios Conforming Array to Differentiable made a crazy cool thing possible: Autograd in Swift.

eaplatanios · 2019-04-23T21:09:02Z

@eaplatanios Conforming Array to Differentiable made a crazy cool thing possible: Autograd in Swift.

Wow this is really cool! On a side note, I also run into the array literal differentiability issue. I was wondering, how can we add a VJP for that?

rxwei · 2019-04-23T21:13:53Z

It requires tweaking the AD transform to recognize _allocateUninitializedArray and a few loads and stores.

Made arrays differentiable.

4a5001c

eaplatanios changed the base branch from master to tensorflow March 8, 2019 19:04

rxwei self-requested a review March 8, 2019 19:34

eaplatanios added 2 commits March 8, 2019 15:29

Addressed Marc's comment.

24f026e

Added documentation for the array differentiable conformance.

18d3154

rxwei reviewed Mar 8, 2019

View reviewed changes

rxwei requested a review from marcrasi March 8, 2019 20:41

rxwei reviewed Mar 8, 2019

View reviewed changes

stdlib/public/core/Array.swift Outdated Show resolved Hide resolved

rxwei added the tensorflow This is for "tensorflow" branch PRs. label Mar 8, 2019

rxwei changed the title ~~[TF] Made arrays differentiable~~ [AutoDiff] [stdlib] Made arrays differentiable Mar 8, 2019

eaplatanios added 3 commits March 8, 2019 16:01

Addressed some of Richard's comments.

4711081

Had to make 'zipLongest' public due to its use in '@inlinable' functi…

d5e1ed9

…ons.

Minor bug fix.

6ec4b8c

rxwei reviewed Mar 8, 2019

View reviewed changes

Addressed Richard's review comments.

0aa8024

marcrasi reviewed Mar 8, 2019

View reviewed changes

stdlib/public/core/Array.swift Outdated Show resolved Hide resolved

marcrasi reviewed Mar 8, 2019

View reviewed changes

stdlib/public/core/Array.swift Outdated Show resolved Hide resolved

dan-zheng reviewed Mar 8, 2019

View reviewed changes

stdlib/public/core/Array.swift Outdated Show resolved Hide resolved

rxwei reviewed Mar 8, 2019

View reviewed changes

stdlib/public/core/Array.swift Outdated Show resolved Hide resolved

stdlib/public/core/Array.swift Outdated Show resolved Hide resolved

stdlib/public/core/Array.swift Outdated Show resolved Hide resolved

stdlib/public/core/Array.swift Outdated Show resolved Hide resolved

eaplatanios added 3 commits March 8, 2019 19:14

Addressed some of the code review feedback.

668890b

Fixed the array differentiable conformance derivation.

37ebb0c

Addressed Marc's comment about 'moved' and 'tangentVector'.

8a278d9

rxwei reviewed Mar 9, 2019

View reviewed changes

stdlib/public/core/Array.swift Outdated Show resolved Hide resolved

stdlib/public/core/Array.swift Outdated Show resolved Hide resolved

Merge branch 'array-differentiable' of https://github.com/eaplatanios…

0346b54

…/swift-language into eaplatanios-array-differentiable

Marc Rasi added 3 commits April 17, 2019 13:35

update to use a single generic type for all array diffble assoc types

2b656d7

Merge branch 'tensorflow' into eaplatanios-array-differentiable

7d98750

add missing comment

d3365ea

rxwei approved these changes Apr 17, 2019

View reviewed changes

address comments

7f8f273

rxwei merged commit 9c34970 into swiftlang:tensorflow Apr 17, 2019

jcatanza mentioned this pull request Apr 23, 2019

v0.3 release. tensorflow/swift#175

Merged

rxwei pushed a commit to rxwei/swift that referenced this pull request May 11, 2019

[AutoDiff] [stdlib] Made arrays differentiable (swiftlang#23183)

2c48150

	public init(_ array: [Element]) { self.base = array }
	public init(_ base: [Element]) { self.base = base }

	/// Construct a view of the given array.
	/// Creates a differentiable view of the given array.

	/// Views the array as the differentiable product manifold of `Element` multiplied with itself
	/// The view of an array as the differentiable product manifold of `Element` multiplied with itself

[AutoDiff] [stdlib] Made arrays differentiable #23183

[AutoDiff] [stdlib] Made arrays differentiable #23183

Uh oh!

Conversation

eaplatanios commented Mar 8, 2019

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

eaplatanios commented Mar 8, 2019

Uh oh!

rxwei commented Mar 8, 2019

Uh oh!

eaplatanios commented Mar 8, 2019

Uh oh!

marcrasi commented Mar 8, 2019

Uh oh!

eaplatanios commented Mar 8, 2019

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

rxwei commented Mar 9, 2019

Uh oh!

eaplatanios commented Mar 9, 2019

Uh oh!

saeta commented Apr 16, 2019

Uh oh!

marcrasi commented Apr 17, 2019

Uh oh!

eaplatanios commented Apr 17, 2019

Uh oh!

rxwei commented Apr 17, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

eaplatanios commented Apr 17, 2019

Uh oh!

marcrasi commented Apr 17, 2019

Uh oh!

marcrasi commented Apr 17, 2019

Uh oh!

eaplatanios commented Apr 17, 2019

Uh oh!

rxwei commented Apr 17, 2019

Uh oh!

marcrasi commented Apr 17, 2019

Uh oh!

marcrasi commented Apr 17, 2019

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rxwei commented Apr 17, 2019 •

edited

Loading