What is the correct way of defining a differentiable function with scalar arguments that returns a vector/array?

58 views Asked by At

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.

0

There are 0 answers