Add @Freezable property wrapper.

`@Freezable` wraps differentiable values and provides toggleable
trainability via the `isFrozen` property.

When `isFrozen` is true, accesses to `value` have a derivative of zero.

Author: dan-zheng

Sources/TensorFlow/Layer.swift

@@ -191,3 +191,34 @@ public final class Parameter<Scalar: TensorFlowScalar> {
         self.value = value
     }
 }
+
+/// A wrapper around a differentiable value with "freezable" derivatives.
+///
+/// When `isFrozen` is true, accesses to `value` have a derivative of zero.
+@_propertyWrapper
+public struct Freezable<Value: Differentiable> : Differentiable {
+    @noDerivative public var isFrozen: Bool = false
+    private var _value: Value
+
+    public init(initialValue: Value) {
+        _value = initialValue
+    }
+
+    @differentiable(vjp: _vjpValue)
+    public var value: Value {
+        get { _value }
+        set { _value = newValue }
+    }
+
+    @usableFromInline
+    func _vjpValue() -> (value: Value, pullback: (Value.TangentVector) -> TangentVector) {
+        return (_value, { [isFrozen = self.isFrozen] v in
+            isFrozen ? .zero : v
+        })
+    }
+
+    public typealias TangentVector = Value.TangentVector
+    public mutating func move(along direction: TangentVector) {
+        _value.move(along: direction)
+    }
+}

Tests/TensorFlowTests/FreezableTests.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 FreezableTests: XCTestCase {
+    func testFreezableParameters() {
+        struct FreezableDense : Layer {
+            @Freezable var weight: Float
+            // Workaround for underlying storage private access.
+            var freezableWeight: Freezable<Float> {
+                get { $weight }
+                set { $weight = newValue }
+            }
+
+            @Freezable var bias: Float
+            // Workaround for underlying storage private access.
+            var freezableBias: Freezable<Float> {
+                get { $bias }
+                set { $bias = newValue }
+            }
+
+            @differentiable
+            func callAsFunction(_ input: Float) -> Float {
+                return input * weight + bias
+            }
+        }
+
+        var dense = FreezableDense(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: `FreezableDense.TangentVector($weight: ..., $bias: ...)
+            // Goal: `FreezableDense.TangentVector(weight: ..., bias: ...)
+            //
+            // XCTAssertEqual(FreezableDense.TangentVector($weight: 4, $bias: 1), gradient.0)
+            XCTAssertEqual(4, gradient.0.$weight)
+            XCTAssertEqual(1, gradient.0.$bias)
+            XCTAssertEqual(2, gradient.1)
+        }
+        dense.freezableWeight.isFrozen = true
+        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 = [
+        ("testFreezableParameters", testFreezableParameters),
+    ]
+}