Skip to content

Uniform variable names #133

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Feb 20, 2021
Merged

Uniform variable names #133

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 24 additions & 24 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,25 +19,25 @@ Suppose we had the function

```julia
fcalls = 0
function f(dx,x) # in-place
function f(y,x) # in-place
global fcalls += 1
for i in 2:length(x)-1
dx[i] = x[i-1] - 2x[i] + x[i+1]
y[i] = x[i-1] - 2x[i] + x[i+1]
end
dx[1] = -2x[1] + x[2]
dx[end] = x[end-1] - 2x[end]
y[1] = -2x[1] + x[2]
y[end] = x[end-1] - 2x[end]
nothing
end

function g(x) # out-of-place
global fcalls += 1
dx = zero(x)
y = zero(x)
for i in 2:length(x)-1
dx[i] = x[i-1] - 2x[i] + x[i+1]
y[i] = x[i-1] - 2x[i] + x[i+1]
end
dx[1] = -2x[1] + x[2]
dx[end] = x[end-1] - 2x[end]
dx
y[1] = -2x[1] + x[2]
y[end] = x[end-1] - 2x[end]
y
end
```

Expand Down Expand Up @@ -85,8 +85,8 @@ forwarddiff_color_jacobian!(jac, f, x, colorvec = colors)
If one only needs to compute products, one can use the operators. For example,

```julia
u = rand(30)
J = JacVec(f,u)
x = rand(30)
J = JacVec(f,x)
```

makes `J` into a matrix-free operator which calculates `J*v` products. For
Expand Down Expand Up @@ -219,14 +219,14 @@ the function signature for Jacobians is `f!(du,u)`, while out-of-place has
The functions for Jacobians are:

```julia
auto_jacvec!(du, f, x, v,
auto_jacvec!(dy, f, x, v,
cache1 = ForwardDiff.Dual{DeivVecTag}.(x, v),
cache2 = ForwardDiff.Dual{DeivVecTag}.(x, v))

auto_jacvec(f, x, v)

# If compute_f0 is false, then `f(cache1,x)` will be computed
num_jacvec!(du,f,x,v,cache1 = similar(v),
num_jacvec!(dy,f,x,v,cache1 = similar(v),
cache2 = similar(v);
compute_f0 = true)
num_jacvec(f,x,v,f0=nothing)
Expand All @@ -235,21 +235,21 @@ num_jacvec(f,x,v,f0=nothing)
For Hessians, the following are provided:

```julia
num_hesvec!(du,f,x,v,
num_hesvec!(dy,f,x,v,
cache1 = similar(v),
cache2 = similar(v),
cache3 = similar(v))

num_hesvec(f,x,v)

numauto_hesvec!(du,f,x,v,
numauto_hesvec!(dy,f,x,v,
cache = ForwardDiff.GradientConfig(f,v),
cache1 = similar(v),
cache2 = similar(v))

numauto_hesvec(f,x,v)

autonum_hesvec!(du,f,x,v,
autonum_hesvec!(dy,f,x,v,
cache1 = similar(v),
cache2 = ForwardDiff.Dual{DeivVecTag}.(x, v),
cache3 = ForwardDiff.Dual{DeivVecTag}.(x, v))
Expand All @@ -258,17 +258,17 @@ autonum_hesvec(f,x,v)
```

In addition,
the following forms allow you to provide a gradient function `g(dx,x)` or `dx=g(x)`
the following forms allow you to provide a gradient function `g(dy,x)` or `dy=g(x)`
respectively:

```julia
num_hesvecgrad!(du,g,x,v,
num_hesvecgrad!(dy,g,x,v,
cache2 = similar(v),
cache3 = similar(v))

num_hesvecgrad(g,x,v)

auto_hesvecgrad!(du,g,x,v,
auto_hesvecgrad!(dy,g,x,v,
cache2 = ForwardDiff.Dual{DeivVecTag}.(x, v),
cache3 = ForwardDiff.Dual{DeivVecTag}.(x, v))

Expand All @@ -287,14 +287,14 @@ optimized these will likely be the fastest.
```julia
using Zygote # Required

numback_hesvec!(du,f,x,v,
numback_hesvec!(dy,f,x,v,
cache1 = similar(v),
cache2 = similar(v))

numback_hesvec(f,x,v)

# Currently errors! See https://github.com/FluxML/Zygote.jl/issues/241
autoback_hesvec!(du,f,x,v,
autoback_hesvec!(dy,f,x,v,
cache2 = ForwardDiff.Dual{DeivVecTag}.(x, v),
cache3 = ForwardDiff.Dual{DeivVecTag}.(x, v))

Expand All @@ -308,9 +308,9 @@ Jacobian-vector and Hessian-vector products where the differentiation takes
place at the vector `u`:

```julia
JacVec(f,u::AbstractArray;autodiff=true)
HesVec(f,u::AbstractArray;autodiff=true)
HesVecGrad(g,u::AbstractArray;autodiff=false)
JacVec(f,x::AbstractArray;autodiff=true)
HesVec(f,x::AbstractArray;autodiff=true)
HesVecGrad(g,x::AbstractArray;autodiff=false)
```

These all have the same interface, where `J*v` utilizes the out-of-place
Expand Down
102 changes: 51 additions & 51 deletions src/differentiation/jaches_products.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,18 +1,18 @@
struct DeivVecTag end

# J(f(x))*v
function auto_jacvec!(du, f, x, v,
function auto_jacvec!(dy, f, x, v,
cache1 = ForwardDiff.Dual{DeivVecTag}.(x, v),
cache2 = ForwardDiff.Dual{DeivVecTag}.(x, v))
cache1 .= Dual{DeivVecTag}.(x, v)
f(cache2,cache1)
du .= partials.(cache2, 1)
dy .= partials.(cache2, 1)
end
function auto_jacvec(f, x, v)
partials.(f(Dual{DeivVecTag}.(x, v)), 1)
end

function num_jacvec!(du,f,x,v,cache1 = similar(v),
function num_jacvec!(dy,f,x,v,cache1 = similar(v),
cache2 = similar(v);
compute_f0 = true)
compute_f0 && (f(cache1,x))
Expand All @@ -22,7 +22,7 @@ function num_jacvec!(du,f,x,v,cache1 = similar(v),
@. x += ϵ*v
f(cache2,x)
@. x -= ϵ*v
@. du = (cache2 - cache1)/ϵ
@. dy = (cache2 - cache1)/ϵ
end

function num_jacvec(f,x,v,f0=nothing)
Expand All @@ -33,7 +33,7 @@ function num_jacvec(f,x,v,f0=nothing)
(f(x.+ϵ.*v) .- _f0)./ϵ
end

function num_hesvec!(du,f,x,v,
function num_hesvec!(dy,f,x,v,
cache1 = similar(v),
cache2 = similar(v),
cache3 = similar(v))
Expand All @@ -48,7 +48,7 @@ function num_hesvec!(du,f,x,v,
g(cache2,x)
@. x -= 2ϵ*v
g(cache3,x)
@. du = (cache2 - cache3)/(2ϵ)
@. dy = (cache2 - cache3)/(2ϵ)
end

function num_hesvec(f,x,v)
Expand All @@ -63,7 +63,7 @@ function num_hesvec(f,x,v)
(gxp - gxm)/(2ϵ)
end

function numauto_hesvec!(du,f,x,v,
function numauto_hesvec!(dy,f,x,v,
cache = ForwardDiff.GradientConfig(f,v),
cache1 = similar(v),
cache2 = similar(v))
Expand All @@ -77,7 +77,7 @@ function numauto_hesvec!(du,f,x,v,
g(cache1,x)
@. x -= 2ϵ*v
g(cache2,x)
@. du = (cache1 - cache2)/(2ϵ)
@. dy = (cache1 - cache2)/(2ϵ)
end

function numauto_hesvec(f,x,v)
Expand All @@ -92,22 +92,22 @@ function numauto_hesvec(f,x,v)
(gxp - gxm)/(2ϵ)
end

function autonum_hesvec!(du,f,x,v,
function autonum_hesvec!(dy,f,x,v,
cache1 = ForwardDiff.Dual{DeivVecTag}.(x, v),
cache2 = ForwardDiff.Dual{DeivVecTag}.(x, v))
cache = FiniteDiff.GradientCache(v[1],cache1,Val{:central})
g = (dx,x) -> FiniteDiff.finite_difference_gradient!(dx,f,x,cache)
cache1 .= Dual{DeivVecTag}.(x, v)
g(cache2,cache1)
du .= partials.(cache2, 1)
dy .= partials.(cache2, 1)
end

function autonum_hesvec(f,x,v)
g = (x) -> FiniteDiff.finite_difference_gradient(f,x)
partials.(g(Dual{DeivVecTag}.(x, v)), 1)
end

function num_hesvecgrad!(du,g,x,v,
function num_hesvecgrad!(dy,g,x,v,
cache2 = similar(v),
cache3 = similar(v))
T = eltype(x)
Expand All @@ -117,7 +117,7 @@ function num_hesvecgrad!(du,g,x,v,
g(cache2,x)
@. x -= 2ϵ*v
g(cache3,x)
@. du = (cache2 - cache3)/(2ϵ)
@. dy = (cache2 - cache3)/(2ϵ)
end

function num_hesvecgrad(g,x,v)
Expand All @@ -131,12 +131,12 @@ function num_hesvecgrad(g,x,v)
(gxp - gxm)/(2ϵ)
end

function auto_hesvecgrad!(du,g,x,v,
function auto_hesvecgrad!(dy,g,x,v,
cache2 = ForwardDiff.Dual{DeivVecTag}.(x, v),
cache3 = ForwardDiff.Dual{DeivVecTag}.(x, v))
cache2 .= Dual{DeivVecTag}.(x, v)
g(cache3,cache2)
du .= partials.(cache3, 1)
dy .= partials.(cache3, 1)
end

function auto_hesvecgrad(g,x,v)
Expand All @@ -145,98 +145,98 @@ end

### Operator Forms

struct JacVec{F,T1,T2,uType}
struct JacVec{F,T1,T2,xType}
f::F
cache1::T1
cache2::T2
u::uType
x::xType
autodiff::Bool
end

function JacVec(f,u::AbstractArray;autodiff=true)
function JacVec(f,x::AbstractArray;autodiff=true)
if autodiff
cache1 = ForwardDiff.Dual{DeivVecTag}.(u, u)
cache2 = ForwardDiff.Dual{DeivVecTag}.(u, u)
cache1 = ForwardDiff.Dual{DeivVecTag}.(x, x)
cache2 = ForwardDiff.Dual{DeivVecTag}.(x, x)
else
cache1 = similar(u)
cache2 = similar(u)
cache1 = similar(x)
cache2 = similar(x)
end
JacVec(f,cache1,cache2,u,autodiff)
JacVec(f,cache1,cache2,x,autodiff)
end

Base.size(L::JacVec) = (length(L.cache1),length(L.cache1))
Base.size(L::JacVec,i::Int) = length(L.cache1)
Base.:*(L::JacVec,x::AbstractVector) = L.autodiff ? auto_jacvec(_u->L.f(_u),L.u,x) : num_jacvec(_u->L.f(_u),L.u,x)
Base.:*(L::JacVec,v::AbstractVector) = L.autodiff ? auto_jacvec(_x->L.f(_x),L.x,v) : num_jacvec(_x->L.f(_x),L.x,v)

function LinearAlgebra.mul!(du::AbstractVector,L::JacVec,v::AbstractVector)
function LinearAlgebra.mul!(dy::AbstractVector,L::JacVec,v::AbstractVector)
if L.autodiff
auto_jacvec!(du,(_du,_u)->L.f(_du,_u),L.u,v,L.cache1,L.cache2)
auto_jacvec!(dy,(_y,_x)->L.f(_y,_x),L.x,v,L.cache1,L.cache2)
else
num_jacvec!(du,(_du,_u)->L.f(_du,_u),L.u,v,L.cache1,L.cache2)
num_jacvec!(dy,(_y,_x)->L.f(_y,_x),L.x,v,L.cache1,L.cache2)
end
end

struct HesVec{F,T1,T2,uType}
struct HesVec{F,T1,T2,xType}
f::F
cache1::T1
cache2::T2
cache3::T2
u::uType
x::xType
autodiff::Bool
end

function HesVec(f,u::AbstractArray;autodiff=true)
function HesVec(f,x::AbstractArray;autodiff=true)
if autodiff
cache1 = ForwardDiff.GradientConfig(f,u)
cache2 = similar(u)
cache3 = similar(u)
cache1 = ForwardDiff.GradientConfig(f,x)
cache2 = similar(x)
cache3 = similar(x)
else
cache1 = similar(u)
cache2 = similar(u)
cache3 = similar(u)
cache1 = similar(x)
cache2 = similar(x)
cache3 = similar(x)
end
HesVec(f,cache1,cache2,cache3,u,autodiff)
HesVec(f,cache1,cache2,cache3,x,autodiff)
end

Base.size(L::HesVec) = (length(L.cache2),length(L.cache2))
Base.size(L::HesVec,i::Int) = length(L.cache2)
Base.:*(L::HesVec,x::AbstractVector) = L.autodiff ? numauto_hesvec(L.f,L.u,x) : num_hesvec(L.f,L.u,x)
Base.:*(L::HesVec,v::AbstractVector) = L.autodiff ? numauto_hesvec(L.f,L.x,v) : num_hesvec(L.f,L.x,v)

function LinearAlgebra.mul!(du::AbstractVector,L::HesVec,v::AbstractVector)
function LinearAlgebra.mul!(dy::AbstractVector,L::HesVec,v::AbstractVector)
if L.autodiff
numauto_hesvec!(du,L.f,L.u,v,L.cache1,L.cache2,L.cache3)
numauto_hesvec!(dy,L.f,L.x,v,L.cache1,L.cache2,L.cache3)
else
num_hesvec!(du,L.f,L.u,v,L.cache1,L.cache2,L.cache3)
num_hesvec!(dy,L.f,L.x,v,L.cache1,L.cache2,L.cache3)
end
end

struct HesVecGrad{G,T1,T2,uType}
g::G
cache1::T1
cache2::T2
u::uType
x::uType
autodiff::Bool
end

function HesVecGrad(g,u::AbstractArray;autodiff=false)
function HesVecGrad(g,x::AbstractArray;autodiff=false)
if autodiff
cache1 = ForwardDiff.Dual{DeivVecTag}.(u, u)
cache2 = ForwardDiff.Dual{DeivVecTag}.(u, u)
cache1 = ForwardDiff.Dual{DeivVecTag}.(x, x)
cache2 = ForwardDiff.Dual{DeivVecTag}.(x, x)
else
cache1 = similar(u)
cache2 = similar(u)
cache1 = similar(x)
cache2 = similar(x)
end
HesVecGrad(g,cache1,cache2,u,autodiff)
HesVecGrad(g,cache1,cache2,x,autodiff)
end

Base.size(L::HesVecGrad) = (length(L.cache2),length(L.cache2))
Base.size(L::HesVecGrad,i::Int) = length(L.cache2)
Base.:*(L::HesVecGrad,x::AbstractVector) = L.autodiff ? auto_hesvecgrad(L.g,L.u,x) : num_hesvecgrad(L.g,L.u,x)
Base.:*(L::HesVecGrad,v::AbstractVector) = L.autodiff ? auto_hesvecgrad(L.g,L.x,v) : num_hesvecgrad(L.g,L.x,v)

function LinearAlgebra.mul!(du::AbstractVector,L::HesVecGrad,v::AbstractVector)
function LinearAlgebra.mul!(dy::AbstractVector,L::HesVecGrad,v::AbstractVector)
if L.autodiff
auto_hesvecgrad!(du,L.g,L.u,v,L.cache1,L.cache2)
auto_hesvecgrad!(dy,L.g,L.x,v,L.cache1,L.cache2)
else
num_hesvecgrad!(du,L.g,L.u,v,L.cache1,L.cache2)
num_hesvecgrad!(dy,L.g,L.x,v,L.cache1,L.cache2)
end
end
Loading