improve differential operator tests

Author: Marc Rasi

stdlib/public/Differentiation/DifferentiationUtilities.swift

@@ -52,6 +52,24 @@ public func differentiableFunction<T, U, R>(
     /*vjp*/ vjp)
 }
 
+/// Create a differentiable function from a vector-Jacobian products function.
+@inlinable
+public func differentiableFunction<T, U, V, R>(
+  from vjp: @escaping (T, U, V)
+           -> (value: R, pullback: (R.TangentVector)
+             -> (T.TangentVector, U.TangentVector, V.TangentVector))
+) -> @differentiable (T, U, V) -> R {
+  Builtin.differentiableFunction_arity3(
+    /*original*/ { vjp($0, $1, $2).value },
+    /*jvp*/ { _, _, _ in
+      fatalError("""
+        Functions formed with `differentiableFunction(from:)` cannot yet \
+        be used with differential-producing differential operators.
+        """)
+    },
+    /*vjp*/ vjp)
+}
+
 /// Returns `x` like an identity function. When used in a context where `x` is
 /// being differentiated with respect to, this function will not produce any 
 /// derivative at `x`.

test/AutoDiff/stdlib/differential_operators.swift

@@ -1,25 +1,77 @@
 // RUN: %empty-directory(%t)
-// RUN: %target-build-swift -enable-experimental-differentiable-programming %s -emit-sil
-// RUN: %target-build-swift -enable-experimental-differentiable-programming %s -o %t/differential_operators
+// RUN: %gyb %s -o %t/differential_operators.swift
+// RUN: %target-build-swift -enable-experimental-differentiable-programming %t/differential_operators.swift -o %t/differential_operators
 // RUN: %target-run %t/differential_operators
 // REQUIRES: executable_test
 
 import _Differentiation
 
 import StdlibUnittest
 
-var DifferentialOperatorTestSuite = TestSuite("Tuple")
+var DifferentialOperatorTestSuite = TestSuite("DifferentialOperator")
 
-DifferentialOperatorTestSuite.test("foo") {
-  expectEqual(1, 1)
+% for arity in range(1, 3 + 1):
+
+% params = ', '.join(['_ x%d: Float' % i for i in range(arity)])
+% pb_return_type = '(' + ', '.join(['Float' for _ in range(arity)]) + ')'
+func exampleVJP_${arity}(${params}) -> (Float, (Float) -> ${pb_return_type}) {
+  (
+    ${' + '.join(['x%d * x%d' % (i, i) for i in range(arity)])},
+    { (${', '.join(['2 * x%d * $0' % i for i in range(arity)])}) }
+  )
+}
+
+% argValues = [i * 10 for i in range(1, arity + 1)]
+% args = ', '.join([str(v) for v in argValues])
+% expectedValue = sum([v * v for v in argValues])
+% expectedGradientValues = [2 * v for v in argValues]
+% expectedGradients = '(' + ', '.join([str(g) for g in expectedGradientValues]) + ')'
+
+DifferentialOperatorTestSuite.test("differentiableFunction_callOriginal_${arity}") {
+  let f = differentiableFunction(from: exampleVJP_${arity})
+  expectEqual(${expectedValue}, f(${args}))
 }
 
-func exampleVJP(_ x: Float) -> (Float, (Float) -> Float) {
-  (x * x, { 2 * $0 })
+DifferentialOperatorTestSuite.test("valueWithPullback_${arity}") {
+  let f = differentiableFunction(from: exampleVJP_${arity})
+  let (value, pb) = valueWithPullback(at: ${args}, in: f)
+  expectEqual(${expectedValue}, value)
+  expectEqual(${expectedGradients}, pb(1))
 }
 
-DifferentialOperatorTestSuite.test("differentiableFunction_callOriginal") {
-  expectEqual(differentiableFunction(from: exampleVJP)(10), 100)
+DifferentialOperatorTestSuite.test("pullback_${arity}") {
+  let f = differentiableFunction(from: exampleVJP_${arity})
+  let pb = pullback(at: ${args}, in: f)
+  expectEqual(${expectedGradients}, pb(1))
 }
 
+DifferentialOperatorTestSuite.test("gradient_${arity}") {
+  let f = differentiableFunction(from: exampleVJP_${arity})
+  let grad = gradient(at: ${args}, in: f)
+  expectEqual(${expectedGradients}, grad)
+}
+
+DifferentialOperatorTestSuite.test("valueWithGradient_${arity}") {
+  let f = differentiableFunction(from: exampleVJP_${arity})
+  let (value, grad) = valueWithGradient(at: ${args}, in: f)
+  expectEqual(${expectedValue}, value)
+  expectEqual(${expectedGradients}, grad)
+}
+
+DifferentialOperatorTestSuite.test("gradient_curried_${arity}") {
+  let f = differentiableFunction(from: exampleVJP_${arity})
+  let gradF = gradient(of: f)
+  expectEqual(${expectedGradients}, gradF(${args}))
+}
+
+DifferentialOperatorTestSuite.test("valueWithGradient_curried_${arity}") {
+  let f = differentiableFunction(from: exampleVJP_${arity})
+  let valueWithGradF = valueWithGradient(of: f)
+  let (value, grad) = valueWithGradF(${args})
+  expectEqual(${expectedValue}, value)
+  expectEqual(${expectedGradients}, grad)
+}
+
+% end
+
 runAllTests()