mirror of
https://github.com/apple/swift.git
synced 2025-12-14 20:36:38 +01:00
89a7663f1d
This adds initial support for differentiation of functions that may produce `Error` result. Essentially we wrap the pullback into `Optional` and emit a diamond-shape control flow pattern depending on whether the pullback value is available or not. VJP emission was modified to accommodate for this. In addition to this, some additional tricks are required as `try_apply` result is not available in the instruction parent block, it is available in normal successor basic block. As a result we can now: - differentiate an active `try_apply` result (that would be produced from `do ... try .. catch` constructions) - `try_apply` when error result is unreachable (usually `try!` and similar source code constructs) - Support (some) throwing functions with builtin differentiation operators. stdlib change will follow. Though we cannot support typed throws here (yet) - Correctly propagate error types during currying around differentiable functions as well as type-checking for `@derivative(of:)` attribute, so we can register custom derivatives for functions producing error result - Added custom derivative for `Optional.??` operator (note that support here is not yet complete as we cannot differentiate through autoclosures, so `x ?? y` works only if `y` is not active, e.g. a constant value). Some fixes here and there
86 lines
2.4 KiB
Swift
86 lines
2.4 KiB
Swift
//===--- OptionalDifferentiation.swift ------------------------*- swift -*-===//
|
|
//
|
|
// This source file is part of the Swift.org open source project
|
|
//
|
|
// Copyright (c) 2020 Apple Inc. and the Swift project authors
|
|
// Licensed under Apache License v2.0 with Runtime Library Exception
|
|
//
|
|
// See https://swift.org/LICENSE.txt for license information
|
|
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
import Swift
|
|
|
|
extension Optional: Differentiable where Wrapped: Differentiable {
|
|
@frozen
|
|
public struct TangentVector: Differentiable, AdditiveArithmetic {
|
|
public typealias TangentVector = Self
|
|
|
|
public var value: Wrapped.TangentVector?
|
|
|
|
@inlinable
|
|
public init(_ value: Wrapped.TangentVector?) {
|
|
self.value = value
|
|
}
|
|
|
|
@inlinable
|
|
public static var zero: Self {
|
|
return Self(.zero)
|
|
}
|
|
|
|
@inlinable
|
|
public static func + (lhs: Self, rhs: Self) -> Self {
|
|
switch (lhs.value, rhs.value) {
|
|
case (nil, nil): return Self(nil)
|
|
case let (x?, nil): return Self(x)
|
|
case let (nil, y?): return Self(y)
|
|
case let (x?, y?): return Self(x + y)
|
|
}
|
|
}
|
|
|
|
@inlinable
|
|
public static func - (lhs: Self, rhs: Self) -> Self {
|
|
switch (lhs.value, rhs.value) {
|
|
case (nil, nil): return Self(nil)
|
|
case let (x?, nil): return Self(x)
|
|
case let (nil, y?): return Self(.zero - y)
|
|
case let (x?, y?): return Self(x - y)
|
|
}
|
|
}
|
|
|
|
@inlinable
|
|
public mutating func move(by offset: TangentVector) {
|
|
if let value = offset.value {
|
|
self.value?.move(by: value)
|
|
}
|
|
}
|
|
}
|
|
|
|
@inlinable
|
|
public mutating func move(by offset: TangentVector) {
|
|
if let value = offset.value {
|
|
self?.move(by: value)
|
|
}
|
|
}
|
|
}
|
|
|
|
extension Optional.TangentVector: CustomReflectable {
|
|
public var customMirror: Mirror {
|
|
return value.customMirror
|
|
}
|
|
}
|
|
|
|
@derivative(of: ??)
|
|
@_transparent
|
|
@_alwaysEmitIntoClient
|
|
func _vjpNilCoalescing<T: Differentiable>(optional: T?, defaultValue: @autoclosure () throws -> T)
|
|
rethrows -> (value: T, pullback: (T.TangentVector) -> Optional<T>.TangentVector) {
|
|
let hasValue = optional != nil
|
|
let value = try optional ?? defaultValue()
|
|
func pullback(_ v: T.TangentVector) -> Optional<T>.TangentVector {
|
|
return hasValue ? .init(v) : .zero
|
|
}
|
|
return (value, pullback)
|
|
}
|