Skip to content

Commit 40ba20b

Browse files
chore: pretend prob is immutable
1 parent 0c249ab commit 40ba20b

File tree

1 file changed

+2
-8
lines changed

1 file changed

+2
-8
lines changed

ext/SciMLBaseZygoteExt.jl

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -304,14 +304,8 @@ end
304304
val = getfield(x, f)
305305
function back(Δ)
306306
Zygote.accum_param(__context__, val, Δ) === nothing && return
307-
if isimmutable(x)
308-
dx = (; Zygote.nt_nothing(x)..., pair(Val(f), Δ, x)...)
309-
(_project(x, dx), nothing)
310-
else
311-
dx = Zygote.grad_mut(__context__, x)
312-
dx[] = (; dx[]..., pair(Val(f), Zygote.accum(getfield(dx[], f), Δ))...)
313-
return (dx[],nothing)
314-
end
307+
dx = (; Zygote.nt_nothing(x)..., pair(Val(f), Δ, x)...)
308+
(_project(x, dx), nothing)
315309
end
316310
Zygote.unwrap(val), back
317311
end

0 commit comments

Comments
 (0)