Skip to content

Commit cc4c332

Browse files
committed
Add Julia fwd, and batch example.
1 parent 9b34cc6 commit cc4c332

File tree

1 file changed

+37
-0
lines changed

1 file changed

+37
-0
lines changed

Diff for: julia_fwd_and_batch/fwd_and_batch.jl

+37
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
using Enzyme
2+
3+
# Defining our test function
4+
function f(x::Array{Float64}, y::Array{Float64})
5+
y[1] = x[1] * x[1] + x[2] * x[1]
6+
return nothing
7+
end;
8+
9+
# To use forward-mode we then have to seed the gradient computation with
10+
x = [2.0, 2.0];
11+
dx = [1.0, 0.0];
12+
y = [0.0];
13+
dy = [0.0];
14+
15+
# Of which the second duplicated element then stores the tangent
16+
Enzyme.autodiff(Forward, f, Duplicated(x, dx), Duplicated(y, dy))
17+
18+
# If we then seek to propagate multiple tangents at the same time to obtain the Hessian in one autodiff call, we then seed
19+
# the following way
20+
y = [0.0];
21+
x = [2.0, 2.0];
22+
23+
vdy = ([0.0],[0.0]);
24+
vdx = ([1.0, 0.0], [0.0, 1.0]);
25+
26+
bx = [0.0, 0.0];
27+
by = [1.0];
28+
vdbx = ([0.0, 0.0], [0.0, 0.0]);
29+
vdby = ([0.0], [0.0]);
30+
31+
# The AD-call then takes the following form
32+
Enzyme.autodiff(
33+
Forward,
34+
(x,y) -> Enzyme.autodiff_deferred(f, x, y),
35+
BatchDuplicated(Duplicated(x, bx), Duplicated.(vdx, vdbx)),
36+
BatchDuplicated(Duplicated(y, by), Duplicated.(vdy, vdby)),
37+
);

0 commit comments

Comments
 (0)