Replace Zygote.@nograd with CRC.@non_differentiable (#230)

* Replace `Zygote.@nograd` with `CRC.@non_differentiable` * Update Project.toml Co-authored-by: Hong Ge <[email protected]>
TuringLang · Aug 23, 2022 · 95c54fa · 95c54fa · devmotion · Aug 23, 2022
1 parent 38c46cf
commit 95c54fa
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 3 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "Bijectors"
 uuid = "76274a88-744f-5084-9051-94815aaf08c4"
-version = "0.10.4"
+version = "0.10.5"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"

diff --git a/src/chainrules.jl b/src/chainrules.jl
@@ -155,3 +155,6 @@ function ChainRulesCore.rrule(::typeof(_transform_inverse_ordered), x::AbstractM
 
     return y, _transform_inverse_ordered_adjoint
 end
+
+# Fixes Zygote's issues with `@debug`
+ChainRulesCore.@non_differentiable _debug(::Any)
diff --git a/src/compat/zygote.jl b/src/compat/zygote.jl
@@ -1,9 +1,8 @@
-using .Zygote: Zygote, @adjoint, @nograd, pullback
+using .Zygote: Zygote, @adjoint, pullback
 
 using Compat: eachcol
 
 @adjoint istraining() = true, _ -> nothing
-@nograd Bijectors._debug
 
 @adjoint function mapvcat(f, args...)
     g(f, args...) = map(f, args...)