derivative_customization.swift

// RUN: %target-run-simple-swift
// REQUIRES: executable_test

import DifferentiationUnittest
import StdlibUnittest

varDerivativeCustomizationTests=TestSuite("DerivativeCustomization")

DerivativeCustomizationTests.testWithLeakChecking("withDerivative"){
do{
varcounter=0
func callback(_ x:inoutTracked<Float>){ counter +=1}
 _ =gradient(at:4){(x:Tracked<Float>)->Tracked<Float>in
 // Non-active value should not be differentiated, so `callback` should
 // not be called.
 _ = x.withDerivative(callback)
return x.withDerivative(callback)+ x.withDerivative(callback)
}
expectEqual(2, counter)
}

expectEqual(
30,
gradient(at:4){(x:Tracked<Float>)in
 x.withDerivative{ $0 =10}+ x.withDerivative{ $0 =20}
})
}

DerivativeCustomizationTests.testWithLeakChecking("withoutDerivative"){
expectEqual(
0,
gradient(at:Tracked<Float>(4)){ x ->Tracked<Float>in
withoutDerivative(at: x){ x in
 x * x * x
}
})

expectEqual(
0,
gradient(at:Tracked<Float>(4)){ x ->Tracked<Float>in
lety=withoutDerivative(at: x)
return y * y * y
})

expectEqual(
2,
gradient(at:Tracked<Float>(4)){ x ->Tracked<Float>in
lety=withoutDerivative(at: x)
return x + y * y * y + x
})
}

runAllTests()