More unthunking in ∇chunk (#180)

* Update utils.jl * version 0.4.5
JuliaML · Jan 4, 2025 · eb5934a · eb5934a · CarloLucibello · Jan 4, 2025
1 parent 204b958
commit eb5934a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "MLUtils"
 uuid = "f1d291b0-491e-4a28-83b9-f70985020b54"
 authors = ["Carlo Lucibello <[email protected]> and contributors"]
-version = "0.4.4"
+version = "0.4.5"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

diff --git a/src/utils.jl b/src/utils.jl
@@ -237,7 +237,8 @@ end
 @non_differentiable _partition_idxs(::Any...)
 
 # Similar to ∇eachslice  https://github.com/JuliaDiff/ChainRules.jl/blob/8108a77a96af5d4b0c460aac393e44f8943f3c5e/src/rulesets/Base/indexing.jl#L77
-function ∇chunk(dys, x, idxs, vd::Val{dim}) where {dim}
+function ∇chunk(dys_raw, x, idxs, vd::Val{dim}) where {dim}
+    dys = unthunk.(unthunk(dys_raw))  # https://github.com/FluxML/Zygote.jl/pull/966#issuecomment-2569227272
     i1 = findfirst(dy -> !(dy isa AbstractZero), dys)
     if i1 === nothing  # all slices are Zero!
         return _zero_fill!(similar(x, float(eltype(x))))