Running out of memory
- Dominant language
- Julia
- Stars
- 586
- Forks
- 108
- Avg merge
- 1d 5h
- Merged PRs (30d)
- 44
Description
Hello,
Thanks for finally supporting mutation in AD!!
I am on a mac and on `julia 1.7` it seems that an example of this size will die (allocates a lot of memory and goes to swap, I guess the kernel kills it). I tried for the function to be non-allocating as below, and didn't help.
```
function fwd3!!(d,c,pp,pc)
n1,n2,n3=size(c)
nt=length(d)
#pp=zeros(n1,n2,n3)
#pc=zeros(n1,n2,n3)
pp.=0.
pc.=0.
pp[n1÷2,n2÷2,n3÷2]=1
for it in 1:nt
for iy in 2:n3-1
for ix in 2:n2-1
for iz in 2:n1-1
pc[iz,ix,iy] = 2pp[iz,ix,iy] - pc[iz,ix,iy] +
1e-8*c[iz,ix,iy]^2*(pp[iz-1,ix,iy]-2pp[iz,ix,iy]+pp[iz+1,ix,iy])+
1e-8*c[iz,ix,iy]^2*(pp[iz,ix-1,iy]-2pp[iz,ix,iy]+pp[iz,ix+1,iy])+
1e-8*c[iz,ix,iy]^2*(pp[iz,ix,iy-1]-2pp[iz,ix,iy]+pp[iz,ix,iy+1])
end
end
end
pc,pp=pp,pc
d[it] = pc[n1÷4,n2÷3,n3÷3]
end
return nothing
end
c=1e3*ones(101,111,105)
pp,pc=similar(c),similar(c)
d=zeros(1500)
@time fwd3!!(d,c,pp,pc)
function adj3!(d,c,dd)
ddc=copy(dd)
dc=zeros(size(c)...)
pp,pc=similar(c),similar(c)
tpp,tpc=zeros(size(c...)),zeros(size(c)...)
Enzyme.autodiff(fwd3!!, Duplicated(d,ddc), Duplicated(c,dc),Duplicated(pp,tpp),Duplicated(pc,tpc))
return dc
end
ddr=rand(size(d)...)
@time dc=adj3!(d,c,ddr); ### and dead
```
Is there a way around that? I understand that the size of the problem is big, but these are the sizes I am interested in for the applications I have in mind :)
Also, while I am here... Is there a way to define pushforwards in `enzyme`? I asked on [discourse](https://discourse.julialang.org/t/pushforward-in-enzyme-jl/83921), so excuse me for asking here as well!
Cheers!
Contributor guide
Assessment
This issue has not been assessed yet.