How to define a differentiable function with scalar arguments that returns a vector/array?

Using Swift
1

I am a beginner in Swift and am testing its AD capabilities. I have the following function

import Foundation
import _Differentiation

@differentiable(reverse)
func numToVec1(_ x: Float) -> Array<Float>.DifferentiableView {
    typealias V = Array<Float>
    return V.DifferentiableView([sin(x), cos(x), x*x])
}

And going by the stdlib source, the derivative of this function is calculated using derivative(at:of:)

let scalarInput: Float = 42.0 
let dnv1 = derivative(at: scalarInput, of: numToVec1)

This compiles, but results in the following runtime error.

_Differentiation/DifferentiationUtilities.swift:40: Fatal error: JVP does not exist. Use '-Xfrontend -enable-experimental-forward-mode-differentiation' to enable differential-first differentiation APIs.
Current stack trace:
0    libswiftCore.so                    0x00007f4107c1aa60 _swift_stdlib_reportFatalErrorInFile + 112
1    libswiftCore.so                    0x00007f410790d3af <unavailable> + 1442735
2    libswiftCore.so                    0x00007f410790d1c7 <unavailable> + 1442247
3    libswiftCore.so                    0x00007f410790bfd0 _assertionFailure(_:_:file:line:flags:) + 364
4    libswift_Differentiation.so        0x00007f410805d464 <unavailable> + 169060
5    output.s                           0x000055abacce9189 <unavailable> + 8585
6    output.s                           0x000055abacce8e37 <unavailable> + 7735
7    libswift_Differentiation.so        0x00007f410805b820 valueWithDifferential<A, B>(at:of:) + 106
8    libswift_Differentiation.so        0x00007f410805bca0 differential<A, B>(at:of:) + 93
9    libswift_Differentiation.so        0x00007f410805c050 derivative<A, B>(at:of:) + 89
10   output.s                           0x000055abacce8c09 <unavailable> + 7177
11   libc.so.6                          0x00007f410759df90 __libc_start_main + 243
12   output.s                           0x000055abacce87fe <unavailable> + 6142
Program terminated with signal: SIGILL

And compiling with the recommended flags results in a compile-time error.

Assertion failed: it != bufferMap.end() && "Tangent buffer should already exist", file D:\a\1\s\swift\lib\SILOptimizer\Differentiation\JVPCloner.cpp, line 317
Please submit a bug report (https://swift.org/contributing/#reporting-bugs) and include the crash backtrace.
Stack dump:
...

Here is a Godbolt link showcasing this problem.
So, before opening an issue I would like to confirm whether I am doing this the intended way or not.

Crossposted on StackOverflow here

2

Forward-mode autodiff is still experimental. It's not nearly as mature as reverse-mode is, so failures are to be expected when attempting to use forward mode (derivative(at:)).

If you actually wanted a (reverse mode) gradient, you can do this (also, DifferentiableView is an implementation detail for most purposes, you can simplify to just [Float]):

@differentiable(reverse)
func numToVec1(_ x: Float) -> [Float] {
    return [sin(x), cos(x), x*x]
}

let valAndPullback = valueWithPullback(at: 4, of: numToVec1)
let gradient = valAndPullback.pullback(([1,1,1]))
print("function output: \(valAndPullback.value), grad: \(gradient))")

prints function output: [-0.7568025, -0.6536436, 16.0], grad: 8.103159)

1 Like