Add @Trainable property wrapper.

`@Trainable` wraps differentiable values and provides toggleable
trainability via the `isTrainable` property.

Author: dan-zheng

Sources/TensorFlow/Layer.swift

@@ -191,3 +191,35 @@ public final class Parameter<Scalar: TensorFlowScalar> {
         self.value = value
     }
 }
+
+/// A wrapper around a differentiable value with toggleable trainability.
+@_propertyWrapper
+public struct Trainable<Value: Differentiable> : Differentiable {
+    @noDerivative public var isTrainable: Bool = true
+    public var _value: Value
+
+    public init(initialValue: Value) {
+        _value = initialValue
+    }
+
+    @differentiable(vjp: _vjpValue)
+    public var value: Value {
+        get { _value.withoutDerivative() }
+        set { _value = newValue }
+    }
+
+    @usableFromInline
+    func _vjpValue() -> (value: Value, pullback: (Value.TangentVector) -> TangentVector) {
+        return (value, { [isTrainable = self.isTrainable] v in
+            if isTrainable {
+                return v
+            }
+            return .zero
+        })
+    }
+
+    public typealias TangentVector = Value.TangentVector
+    public mutating func move(along direction: TangentVector) {
+        _value.move(along: direction)
+    }
+}

Tests/TensorFlowTests/TrainableTests.swift

@@ -0,0 +1,74 @@
+// Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import XCTest
+@testable import TensorFlow
+
+final class TrainableTests: XCTestCase {
+    func testTrainableLayer() {
+        struct TrainableDense : Layer {
+            @Trainable var weight: Float
+            // Workaround for internal storage property.
+            var trainableWeight: Trainable<Float> {
+                get { $weight }
+                set { $weight = newValue }
+            }
+
+            @Trainable var bias: Float
+            // Workaround for internal storage property.
+            var trainableBias: Trainable<Float> {
+                get { $bias }
+                set { $bias = newValue }
+            }
+
+            @differentiable
+            func callAsFunction(_ input: Float) -> Float {
+                return input * weight + bias
+            }
+        }
+
+        var dense = TrainableDense(weight: 2, bias: 3)
+        let x: Float = 4
+        do {
+            let (value, gradient) = valueWithGradient(at: dense, x) { dense, x in dense(x) }
+            XCTAssertEqual(11, value)
+            // FIXME: '$' is not a valid identifier:
+            //     cannot declare entity named '$weight'; the '$' prefix is reserved for
+            //     implicitly-synthesized declarations.
+            //
+            // Tentative solution: change `Differentiable` derived conformances to use original
+            // property names instead of underlying storage property names.
+            //
+            // Now: `TrainableDense.TangentVector($weight: ..., $bias: ...)
+            // Goal: `TrainableDense.TangentVector(weight: ..., bias: ...)
+            //
+            // XCTAssertEqual(TrainableDense.TangentVector($weight: 4, $bias: 1), gradient.0)
+            XCTAssertEqual(4, gradient.0.$weight)
+            XCTAssertEqual(1, gradient.0.$bias)
+            XCTAssertEqual(2, gradient.1)
+        }
+        dense.trainableWeight.isTrainable = false
+        do {
+            let (value, gradient) = valueWithGradient(at: dense, x) { dense, x in dense(x) }
+            XCTAssertEqual(11, value)
+            XCTAssertEqual(0, gradient.0.$weight)
+            XCTAssertEqual(1, gradient.0.$bias)
+            XCTAssertEqual(2, gradient.1)
+        }
+    }
+
+    static var allTests = [
+        ("testTrainableLayer", testTrainableLayer),
+    ]
+}