Expand unthunk-tangent with more methods (#28)

* Expand unthunk-tangent to more methods * Fix * Bump to 0.2.6
FluxML · Jan 20, 2025 · a4d77ab · a4d77ab · ToucheSir · Jan 20, 2025
1 parent f9bf0e3
commit a4d77ab
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "ZygoteRules"
 uuid = "700de1a5-db45-46bc-99cf-38207098b444"
-version = "0.2.5"
+version = "0.2.6"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

diff --git a/src/adjoint.jl b/src/adjoint.jl
@@ -1,6 +1,6 @@
 using MacroTools
 using MacroTools: @q, combinedef
-using ChainRulesCore: AbstractZero
+using ChainRulesCore: AbstractZero, AbstractThunk, @non_differentiable
 
 function named(arg)
   if isexpr(arg, :(::)) && length(arg.args) == 1
@@ -37,6 +37,12 @@ function unthunk_tangent end
 @inline unthunk_tangent(x) = x
 @inline unthunk_tangent(x::Tuple) = map(unthunk_tangent, x)
 @inline unthunk_tangent(x::NamedTuple) = map(unthunk_tangent, x)
+@inline unthunk_tangent(x::AbstractThunk) = wrap_chainrules_output(unthunk(x))
+@inline unthunk_tangent(x::NTuple{N,<:Number}) where N = x
+@inline unthunk_tangent(x::AbstractArray{<:Number,N}) where N = x
+@inline unthunk_tangent(x::AbstractArray) = map(unthunk_tangent, x)
+unthunk_tangent(d::IdDict) = IdDict([unthunk_tangent(k) => unthunk_tangent(v) for (k, v) in d])
+@non_differentiable unthunk_tangent(::IdDict)
 
 
 function gradm(ex, mut = false, keepthunks = false)