TastyPancakes/Flux.jl
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #413 from mcabbott/patch-2

Browse Source 
evaluate both 2-ary DiffRules only when needed
This commit is contained in:
Mike J Innes 2018-10-03 12:02:12 +01:00 committed by GitHub
commit 683bbec71c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
src/tracker/scalar.jl
View File

@ -63,7 +63,9 @@ for (M, f, arity) in DiffRules.diffrules()
da, db = DiffRules.diffrule(M, f, :a, :b)
f = :($M.$f)
@eval begin
@grad $f(a::Real, b::Real) = $f(data(a), data(b)), Δ -> (Δ * $da, Δ * $db)
@grad $f(a::TrackedReal, b::TrackedReal) = $f(data(a), data(b)), Δ -> (Δ * $da, Δ * $db)
@grad $f(a::TrackedReal, b::Real) = $f(data(a), b), Δ -> (Δ * $da, zero(b))
@grad $f(a::Real, b::TrackedReal) = $f(a, data(b)), Δ -> (zero(a), Δ * $db)
$f(a::TrackedReal, b::TrackedReal) = track($f, a, b)
$f(a::TrackedReal, b::Real) = track($f, a, b)
$f(a::Real, b::TrackedReal) = track($f, a, b)