Skip to content

Commit

Permalink
Type-stable jacobian & gradient for unprepared ForwardDiff (#541)
Browse files Browse the repository at this point in the history 
* Make ForwardDiff's unprepared jacobian and gradient type-stable with fixed chunk size

* Avoid ambiguity
  • Loading branch information
@gdalle
gdalle authored Oct 4, 2024
1 parent c47e26c commit 84378d7
Show file tree
Hide file tree
Showing 3 changed files with 130 additions and 66 deletions.
2 changes: 1 addition & 1 deletion DifferentiationInterface/Project.toml
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "DifferentiationInterface"
uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
authors = ["Guillaume Dalle", "Adrian Hill"]
version = "0.6.5"
version = "0.6.6"

[deps]
ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
Expand Down
130 changes: 87 additions & 43 deletions DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -159,38 +159,60 @@ end

## Gradient

### Unprepared
### Unprepared, only when chunk size not specified

function DI.value_and_gradient!(
f::F, grad, ::AutoForwardDiff, x, contexts::Vararg{Context,C}
) where {F,C}
fc = with_contexts(f, contexts...)
result = DiffResult(zero(eltype(x)), (grad,))
result = gradient!(result, fc, x)
y = DR.value(result)
grad === DR.gradient(result) || copyto!(grad, DR.gradient(result))
return y, grad
f::F, grad, backend::AutoForwardDiff{chunksize}, x, contexts::Vararg{Context,C}
) where {F,C,chunksize}
if isnothing(chunksize)
fc = with_contexts(f, contexts...)
result = DiffResult(zero(eltype(x)), (grad,))
result = gradient!(result, fc, x)
y = DR.value(result)
grad === DR.gradient(result) || copyto!(grad, DR.gradient(result))
return y, grad
else
prep = DI.prepare_gradient(f, backend, x, contexts...)
return DI.value_and_gradient!(f, grad, prep, backend, x, contexts...)
end
end

function DI.value_and_gradient(
f::F, ::AutoForwardDiff, x, contexts::Vararg{Context,C}
) where {F,C}
fc = with_contexts(f, contexts...)
result = GradientResult(x)
result = gradient!(result, fc, x)
return DR.value(result), DR.gradient(result)
f::F, backend::AutoForwardDiff{chunksize}, x, contexts::Vararg{Context,C}
) where {F,C,chunksize}
if isnothing(chunksize)
fc = with_contexts(f, contexts...)
result = GradientResult(x)
result = gradient!(result, fc, x)
return DR.value(result), DR.gradient(result)
else
prep = DI.prepare_gradient(f, backend, x, contexts...)
return DI.value_and_gradient(f, prep, backend, x, contexts...)
end
end

function DI.gradient!(
f::F, grad, ::AutoForwardDiff, x, contexts::Vararg{Context,C}
) where {F,C}
fc = with_contexts(f, contexts...)
return gradient!(grad, fc, x)
f::F, grad, backend::AutoForwardDiff{chunksize}, x, contexts::Vararg{Context,C}
) where {F,C,chunksize}
if isnothing(chunksize)
fc = with_contexts(f, contexts...)
return gradient!(grad, fc, x)
else
prep = DI.prepare_gradient(f, backend, x, contexts...)
return DI.gradient!(f, grad, prep, backend, x, contexts...)
end
end

function DI.gradient(f::F, ::AutoForwardDiff, x, contexts::Vararg{Context,C}) where {F,C}
fc = with_contexts(f, contexts...)
return gradient(fc, x)
function DI.gradient(
f::F, backend::AutoForwardDiff{chunksize}, x, contexts::Vararg{Context,C}
) where {F,C,chunksize}
if isnothing(chunksize)
fc = with_contexts(f, contexts...)
return gradient(fc, x)
else
prep = DI.prepare_gradient(f, backend, x, contexts...)
return DI.gradient(f, prep, backend, x, contexts...)
end
end

### Prepared
Expand Down Expand Up @@ -252,37 +274,59 @@ end

## Jacobian

### Unprepared
### Unprepared, only when chunk size not specified

function DI.value_and_jacobian!(
f::F, jac, ::AutoForwardDiff, x, contexts::Vararg{Context,C}
) where {F,C}
fc = with_contexts(f, contexts...)
y = fc(x)
result = DiffResult(y, (jac,))
result = jacobian!(result, fc, x)
y = DR.value(result)
jac === DR.jacobian(result) || copyto!(jac, DR.jacobian(result))
return y, jac
f::F, jac, backend::AutoForwardDiff{chunksize}, x, contexts::Vararg{Context,C}
) where {F,C,chunksize}
if isnothing(chunksize)
fc = with_contexts(f, contexts...)
y = fc(x)
result = DiffResult(y, (jac,))
result = jacobian!(result, fc, x)
y = DR.value(result)
jac === DR.jacobian(result) || copyto!(jac, DR.jacobian(result))
return y, jac
else
prep = DI.prepare_jacobian(f, backend, x, contexts...)
return DI.value_and_jacobian!(f, jac, prep, backend, x, contexts...)
end
end

function DI.value_and_jacobian(
f::F, ::AutoForwardDiff, x, contexts::Vararg{Context,C}
) where {F,C}
fc = with_contexts(f, contexts...)
return fc(x), jacobian(fc, x)
f::F, backend::AutoForwardDiff{chunksize}, x, contexts::Vararg{Context,C}
) where {F,C,chunksize}
if isnothing(chunksize)
fc = with_contexts(f, contexts...)
return fc(x), jacobian(fc, x)
else
prep = DI.prepare_jacobian(f, backend, x, contexts...)
return DI.value_and_jacobian(f, prep, backend, x, contexts...)
end
end

function DI.jacobian!(
f::F, jac, ::AutoForwardDiff, x, contexts::Vararg{Context,C}
) where {F,C}
fc = with_contexts(f, contexts...)
return jacobian!(jac, fc, x)
f::F, jac, backend::AutoForwardDiff{chunksize}, x, contexts::Vararg{Context,C}
) where {F,C,chunksize}
if isnothing(chunksize)
fc = with_contexts(f, contexts...)
return jacobian!(jac, fc, x)
else
prep = DI.prepare_jacobian(f, backend, x, contexts...)
return DI.jacobian!(f, jac, prep, backend, x, contexts...)
end
end

function DI.jacobian(f::F, ::AutoForwardDiff, x, contexts::Vararg{Context,C}) where {F,C}
fc = with_contexts(f, contexts...)
return jacobian(fc, x)
function DI.jacobian(
f::F, backend::AutoForwardDiff{chunksize}, x, contexts::Vararg{Context,C}
) where {F,C,chunksize}
if isnothing(chunksize)
fc = with_contexts(f, contexts...)
return jacobian(fc, x)
else
prep = DI.prepare_jacobian(f, backend, x, contexts...)
return DI.jacobian(f, prep, backend, x, contexts...)
end
end

### Prepared
Expand Down
64 changes: 42 additions & 22 deletions DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/twoarg.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -209,39 +209,59 @@ end

## Jacobian

### Unprepared
### Unprepared, only when chunk size is not specified

function DI.value_and_jacobian(
f!::F, y, ::AutoForwardDiff, x, contexts::Vararg{Context,C}
) where {F,C}
fc! = with_contexts(f!, contexts...)
jac = similar(y, length(y), length(x))
result = MutableDiffResult(y, (jac,))
result = jacobian!(result, fc!, y, x)
return DiffResults.value(result), DiffResults.jacobian(result)
f!::F, y, backend::AutoForwardDiff{chunksize}, x, contexts::Vararg{Context,C}
) where {F,C,chunksize}
if isnothing(chunksize)
fc! = with_contexts(f!, contexts...)
jac = similar(y, length(y), length(x))
result = MutableDiffResult(y, (jac,))
result = jacobian!(result, fc!, y, x)
return DiffResults.value(result), DiffResults.jacobian(result)
else
prep = DI.prepare_jacobian(f!, y, backend, x, contexts...)
return DI.value_and_jacobian(f!, y, prep, backend, x, contexts...)
end
end

function DI.value_and_jacobian!(
f!::F, y, jac, ::AutoForwardDiff, x, contexts::Vararg{Context,C}
) where {F,C}
fc! = with_contexts(f!, contexts...)
result = MutableDiffResult(y, (jac,))
result = jacobian!(result, fc!, y, x)
return DiffResults.value(result), DiffResults.jacobian(result)
f!::F, y, jac, backend::AutoForwardDiff{chunksize}, x, contexts::Vararg{Context,C}
) where {F,C,chunksize}
if isnothing(chunksize)
fc! = with_contexts(f!, contexts...)
result = MutableDiffResult(y, (jac,))
result = jacobian!(result, fc!, y, x)
return DiffResults.value(result), DiffResults.jacobian(result)
else
prep = DI.prepare_jacobian(f!, y, backend, x, contexts...)
return DI.value_and_jacobian!(f!, y, jac, prep, backend, x, contexts...)
end
end

function DI.jacobian(
f!::F, y, ::AutoForwardDiff, x, contexts::Vararg{Context,C}
) where {F,C}
fc! = with_contexts(f!, contexts...)
return jacobian(fc!, y, x)
f!::F, y, backend::AutoForwardDiff{chunksize}, x, contexts::Vararg{Context,C}
) where {F,C,chunksize}
if isnothing(chunksize)
fc! = with_contexts(f!, contexts...)
return jacobian(fc!, y, x)
else
prep = DI.prepare_jacobian(f!, y, backend, x, contexts...)
return DI.jacobian(f!, y, prep, backend, x, contexts...)
end
end

function DI.jacobian!(
f!::F, y, jac, ::AutoForwardDiff, x, contexts::Vararg{Context,C}
) where {F,C}
fc! = with_contexts(f!, contexts...)
return jacobian!(jac, fc!, y, x)
f!::F, y, jac, backend::AutoForwardDiff{chunksize}, x, contexts::Vararg{Context,C}
) where {F,C,chunksize}
if isnothing(chunksize)
fc! = with_contexts(f!, contexts...)
return jacobian!(jac, fc!, y, x)
else
prep = DI.prepare_jacobian(f!, y, backend, x, contexts...)
return DI.jacobian!(f!, y, jac, prep, backend, x, contexts...)
end
end

### Prepared
Expand Down

0 comments on commit 84378d7

Please sign in to comment.