Skip to content

Commit

Permalink
test: enable runtime activity for now
Browse files Browse the repository at this point in the history
  • Loading branch information
avik-pal committed Sep 26, 2024
1 parent ef14bab commit f5aefbe
Showing 1 changed file with 17 additions and 25 deletions.
42 changes: 17 additions & 25 deletions test/enzyme.jl
Original file line number Diff line number Diff line change
Expand Up @@ -177,47 +177,39 @@ Enzyme.autodiff(Reverse, f3, Duplicated(copy(A), dA), Duplicated(copy(b1), db1),
A = rand(n, n);
dA = zeros(n, n);
b1 = rand(n);
for alg in (

function fnice(A, b, alg)
prob = LinearProblem(A, b)
sol1 = solve(prob, alg)
return sum(sol1.u)
end

@testset for alg in (
LUFactorization(),
RFLUFactorization() # KrylovJL_GMRES(), fails
)
@show alg
function fb(b)
prob = LinearProblem(A, b)

sol1 = solve(prob, alg)
fb_closure = b -> fnice(A, b, alg)

sum(sol1.u)
end
fb(b1)

fd_jac = FiniteDiff.finite_difference_jacobian(fb, b1) |> vec
fd_jac = FiniteDiff.finite_difference_jacobian(fb_closure, b1) |> vec
@show fd_jac

en_jac = map(onehot(b1)) do db1
eres = Enzyme.autodiff(Forward, fb, Duplicated(copy(b1), db1))
eres[1]
return only(Enzyme.autodiff(set_runtime_activity(Forward), fnice,
Const(A), Duplicated(b1, db1), Const(alg)))
end |> collect
@show en_jac

@test en_jacfd_jac rtol=1e-4

function fA(A)
prob = LinearProblem(A, b1)

sol1 = solve(prob, alg)
fA_closure = A -> fnice(A, b1, alg)

sum(sol1.u)
end
fA(A)

fd_jac = FiniteDiff.finite_difference_jacobian(fA, A) |> vec
fd_jac = FiniteDiff.finite_difference_jacobian(fA_closure, A) |> vec
@show fd_jac

en_jac = map(onehot(A)) do dA
eres = Enzyme.autodiff(Forward, fA, Duplicated(copy(A), dA))
eres[1]
end |> collect
return only(Enzyme.autodiff(set_runtime_activity(Forward), fnice,
Duplicated(A, dA), Const(b1), Const(alg)))
end |> collect |> (x -> reshape(x, n, n))
@show en_jac

@test en_jacfd_jac rtol=1e-4
Expand Down

0 comments on commit f5aefbe

Please sign in to comment.