Files
swift-mirror/test/AutoDiff/Sema/DerivativeRegistrationCrossFile/Inputs/derivatives.swift
Anton Korobeynikov 6e2c4faa34 [AutoDiff] Lookup for custom derivatives in non-primary source files (#58965)
* Lookup for custom derivatives in non-primary source files after typecheck is finished for the primary source.

This registers all custom derivatives before autodiff transformations and makes them available to them.

Fully resolves #55170
2022-07-18 11:52:02 -07:00

36 lines
901 B
Swift

import _Differentiation
@inlinable
@derivative(of: min)
func minVJP<T: Comparable & Differentiable>(
_ x: T,
_ y: T
) -> (value: T, pullback: (T.TangentVector) -> (T.TangentVector, T.TangentVector)) {
func pullback(_ v: T.TangentVector) -> (T.TangentVector, T.TangentVector) {
if x <= y {
return (v, .zero)
}
else {
return (.zero, v)
}
}
return (value: min(x, y), pullback: pullback)
}
@inlinable
@derivative(of: max)
func maxVJP<T: Comparable & Differentiable>(
_ x: T,
_ y: T
) -> (value: T, pullback: (T.TangentVector) -> (T.TangentVector, T.TangentVector)) {
func pullback(_ v: T.TangentVector) -> (T.TangentVector, T.TangentVector) {
if x < y {
return (.zero, v)
}
else {
return (v, .zero)
}
}
return (value: max(x, y), pullback: pullback)
}