Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add @non_differentiable rules for basic random functions #2153

Open
prbzrg opened this issue Nov 6, 2023 · 2 comments
Open

Add @non_differentiable rules for basic random functions #2153

prbzrg opened this issue Nov 6, 2023 · 2 comments
Labels
enhancement New feature or request good first issue Good for newcomers

Comments

@prbzrg
Copy link

prbzrg commented Nov 6, 2023

I need to have these lines in my codes to make AD works.

@non_differentiable CUDA.zeros(::Any...)
@non_differentiable CUDA.rand(::Any...)
@non_differentiable CUDA.randn(::Any...)

Can we add them to the extension? Adapting from
https://github.com/JuliaDiff/ChainRules.jl/blob/main/src/rulesets/Random/random.jl
@prbzrg prbzrg added the bug Something isn't working label Nov 6, 2023
@prbzrg
Copy link
Author

prbzrg commented Nov 6, 2023

MRE:

julia> using Zygote, CUDA

julia> f(x) = sum(CUDA.rand()*x)
f (generic function with 1 method)

julia> Zygote.gradient(f, ones(10))
ERROR: Compiling Tuple{typeof(lock), CUDA.APIUtils.var"#9#12"{CUDA.APIUtils.var"#8#11", CUDA.APIUtils.HandleCache{CuContext, CUDA.CURAND.RNG}, CuContext}, ReentrantLock}: try/catch is not supported.
Refer to the Zygote documentation for fixes.
https://fluxml.ai/Zygote.jl/latest/limitations

Stacktrace:
  [1] macro expansion
    @ C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:101 [inlined]
  [2] _pullback(::Zygote.Context{false}, ::typeof(lock), ::CUDA.APIUtils.var"#9#12"{CUDA.APIUtils.var"#8#11", CUDA.APIUtils.HandleCache{CuContext, CUDA.CURAND.RNG}, CuContext}, ::ReentrantLock)
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:101
  [3] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\lib\utils\cache.jl:26 [inlined]
  [4] _pullback(ctx::Zygote.Context{false}, f::CUDA.APIUtils.var"#check_cache#10"{CUDA.APIUtils.HandleCache{CuContext, CUDA.CURAND.RNG}, CuContext}, args::CUDA.APIUtils.var"#8#11")
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
  [5] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\lib\utils\cache.jl:26 [inlined]
  [6] _pullback(::Zygote.Context{false}, ::CUDA.APIUtils.var"#check_cache#10"{CUDA.APIUtils.HandleCache{CuContext, CUDA.CURAND.RNG}, CuContext})
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
  [7] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\lib\utils\cache.jl:41 [inlined]
  [8] _pullback(::Zygote.Context{false}, ::typeof(pop!), ::CUDA.CURAND.var"#107#113", ::CUDA.APIUtils.HandleCache{CuContext, CUDA.CURAND.RNG}, ::CuContext)
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
  [9] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\lib\curand\CURAND.jl:38 [inlined]
 [10] _pullback(ctx::Zygote.Context{false}, f::CUDA.CURAND.var"#new_state#112", args::NamedTuple{(:device, :context, :stream, :math_mode, :math_precision), Tuple{CuDevice, CuContext, CuStream, CUDA.MathMode, Symbol}})
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
 [11] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\lib\curand\CURAND.jl:52 [inlined]
 [12] _pullback(::Zygote.Context{false}, ::CUDA.CURAND.var"#110#116"{CUDA.CURAND.var"#new_state#112", NamedTuple{(:device, :context, :stream, :math_mode, :math_precision), Tuple{CuDevice, CuContext, CuStream, CUDA.MathMode, Symbol}}})
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
 [13] _pullback
    @ .\dict.jl:468 [inlined]
 [14] _pullback(::Zygote.Context{false}, ::typeof(get!), ::CUDA.CURAND.var"#110#116"{CUDA.CURAND.var"#new_state#112", NamedTuple{(:device, :context, :stream, :math_mode, :math_precision), Tuple{CuDevice, CuContext, CuStream, CUDA.MathMode, Symbol}}}, ::Dict{CuContext, NamedTuple{(:rng,), Tuple{CUDA.CURAND.RNG}}}, ::CuContext)
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
 [15] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\lib\curand\CURAND.jl:51 [inlined]
 [16] _pullback(::Zygote.Context{false}, ::typeof(CUDA.CURAND.default_rng))
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
 [17] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\src\random.jl:267 [inlined]
 [18] _pullback(::Zygote.Context{false}, ::typeof(CUDA.curand_rng))
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
 [19] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\src\random.jl:287 [inlined]
 [20] _pullback(::Zygote.Context{false}, ::typeof(CUDA.rand), ::Type{Float32}, ::Int64)
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
 [21] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\src\random.jl:331 [inlined]
 [22] _pullback(ctx::Zygote.Context{false}, f::typeof(CUDA.rand), args::Type{Float32})
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
 [23] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\src\random.jl:331 [inlined]
 [24] _pullback(::Zygote.Context{false}, ::typeof(CUDA.rand))
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
 [25] _pullback
    @ .\REPL[10]:1 [inlined]
 [26] _pullback(ctx::Zygote.Context{false}, f::typeof(f), args::Vector{Float64})
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
 [27] pullback(f::Function, cx::Zygote.Context{false}, args::Vector{Float64})
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface.jl:44
 [28] pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface.jl:42 [inlined]
 [29] gradient(f::Function, args::Vector{Float64})
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface.jl:96
 [30] top-level scope
    @ REPL[11]:1
 [31] top-level scope
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\src\initialization.jl:205

julia> f2(x) = sum(rand()*x)
f2 (generic function with 1 method)

julia> Zygote.gradient(f2, ones(10))
(Fill(0.9066500056252977, 10),)

@prbzrg
Copy link
Author

prbzrg commented Nov 6, 2023 •
edited
Loading 

julia> f3(x) = sum(CUDA.ones(10) .* x)
f3 (generic function with 1 method)

julia> Zygote.gradient(f3, ones(10))
ERROR: `llvmcall` must be compiled to be called
Stacktrace:
  [1] macro expansion
    @ C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:101 [inlined]
  [2] _pullback(::Zygote.Context{false}, ::Core.IntrinsicFunction, ::String, ::Type{Int64}, ::Type{Tuple{Ptr{Int64}}}, ::Ptr{Int64})
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:101
  [3] _pullback
    @ .\atomics.jl:358 [inlined]
  [4] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\lib\utils\threading.jl:27 [inlined]
  [5] macro expansion
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\lib\utils\memoization.jl:69 [inlined]
  [6] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\src\pool.jl:163 [inlined]
  [7] _pullback(ctx::Zygote.Context{false}, f::typeof(CUDA.stream_ordered), args::CuDevice)
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
  [8] macro expansion
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\src\pool.jl:432 [inlined]
  [9] macro expansion
    @ .\timing.jl:393 [inlined]
 [10] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\src\pool.jl:431 [inlined]
 [11] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\src\pool.jl:427 [inlined]
 [12] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\src\pool.jl:417 [inlined]
 [13] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\src\pool.jl:411 [inlined]
 [14] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\src\array.jl:74 [inlined]
 [15] _pullback(::Zygote.Context{false}, ::Type{CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, ::UndefInitializer, ::Tuple{Int64})
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
 [16] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\src\array.jl:136 [inlined]
 [17] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\src\array.jl:155 [inlined]
 [18] _pullback(::Zygote.Context{false}, ::Type{CuArray{Float32}}, ::UndefInitializer, ::Int64)
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
 [19] _apply(::Function, ::Vararg{Any})
    @ Core .\boot.jl:838
 [20] adjoint
    @ C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\lib\lib.jl:203 [inlined]
 [21] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\ZygoteRules\4nXuu\src\adjoint.jl:66 [inlined]
 [22] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\src\array.jl:632 [inlined]
 [23] _pullback(::Zygote.Context{false}, ::typeof(CUDA.ones), ::Type{Float32}, ::Int64)
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
 [24] _apply(::Function, ::Vararg{Any})
    @ Core .\boot.jl:838
 [25] adjoint
    @ C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\lib\lib.jl:203 [inlined]
 [26] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\ZygoteRules\4nXuu\src\adjoint.jl:66 [inlined]
 [27] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\src\array.jl:634 [inlined]
 [28] _pullback(ctx::Zygote.Context{false}, f::typeof(CUDA.ones), args::Int64)
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
 [29] _pullback
    @ .\REPL[14]:1 [inlined]
 [30] _pullback(ctx::Zygote.Context{false}, f::typeof(f3), args::Vector{Float64})
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
 [31] pullback(f::Function, cx::Zygote.Context{false}, args::Vector{Float64})
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface.jl:44
 [32] pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface.jl:42 [inlined]
 [33] gradient(f::Function, args::Vector{Float64})
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface.jl:96
 [34] top-level scope
    @ REPL[15]:1
 [35] top-level scope
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\src\initialization.jl:205

julia> f4(x) = sum(ones(10) .* x)
f4 (generic function with 1 method)

julia> Zygote.gradient(f4, ones(10))
([1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],)

julia> Zygote.gradient(f3, CUDA.ones(10))
ERROR: Compiling Tuple{typeof(lock), CUDA.APIUtils.var"#9#12"{CUDA.APIUtils.var"#8#11", CUDA.APIUtils.HandleCache{CuContext, CUDA.CURAND.RNG}, CuContext}, ReentrantLock}: try/catch is not supported.
Refer to the Zygote documentation for fixes.
https://fluxml.ai/Zygote.jl/latest/limitations

Stacktrace:
  [1] macro expansion
    @ C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:101 [inlined]
  [2] _pullback(::Zygote.Context{false}, ::typeof(lock), ::CUDA.APIUtils.var"#9#12"{CUDA.APIUtils.var"#8#11", CUDA.APIUtils.HandleCache{CuContext, CUDA.CURAND.RNG}, CuContext}, ::ReentrantLock)
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:101
  [3] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\lib\utils\cache.jl:26 [inlined]
  [4] _pullback(ctx::Zygote.Context{false}, f::CUDA.APIUtils.var"#check_cache#10"{CUDA.APIUtils.HandleCache{CuContext, CUDA.CURAND.RNG}, CuContext}, args::CUDA.APIUtils.var"#8#11")
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
  [5] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\lib\utils\cache.jl:26 [inlined]
  [6] _pullback(::Zygote.Context{false}, ::CUDA.APIUtils.var"#check_cache#10"{CUDA.APIUtils.HandleCache{CuContext, CUDA.CURAND.RNG}, CuContext})
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
  [7] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\lib\utils\cache.jl:41 [inlined]
  [8] _pullback(::Zygote.Context{false}, ::typeof(pop!), ::CUDA.CURAND.var"#107#113", ::CUDA.APIUtils.HandleCache{CuContext, CUDA.CURAND.RNG}, ::CuContext)
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
  [9] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\lib\curand\CURAND.jl:38 [inlined]
 [10] _pullback(ctx::Zygote.Context{false}, f::CUDA.CURAND.var"#new_state#112", args::NamedTuple{(:device, :context, :stream, :math_mode, :math_precision), Tuple{CuDevice, CuContext, CuStream, CUDA.MathMode, Symbol}})
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
 [11] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\lib\curand\CURAND.jl:52 [inlined]
 [12] _pullback(::Zygote.Context{false}, ::CUDA.CURAND.var"#110#116"{CUDA.CURAND.var"#new_state#112", NamedTuple{(:device, :context, :stream, :math_mode, :math_precision), Tuple{CuDevice, CuContext, CuStream, CUDA.MathMode, Symbol}}})
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
 [13] _pullback
    @ .\dict.jl:468 [inlined]
 [14] _pullback(::Zygote.Context{false}, ::typeof(get!), ::CUDA.CURAND.var"#110#116"{CUDA.CURAND.var"#new_state#112", NamedTuple{(:device, :context, :stream, :math_mode, :math_precision), Tuple{CuDevice, CuContext, CuStream, CUDA.MathMode, Symbol}}}, ::Dict{CuContext, NamedTuple{(:rng,), Tuple{CUDA.CURAND.RNG}}}, ::CuContext)
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
 [15] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\lib\curand\CURAND.jl:51 [inlined]
 [16] _pullback(::Zygote.Context{false}, ::typeof(CUDA.CURAND.default_rng))
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
 [17] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\src\random.jl:267 [inlined]
 [18] _pullback(::Zygote.Context{false}, ::typeof(CUDA.curand_rng))
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
 [19] _pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\src\random.jl:321 [inlined]
 [20] _pullback(ctx::Zygote.Context{false}, f::typeof(CUDA.rand), args::Int64)
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
 [21] _pullback
    @ .\REPL[16]:1 [inlined]
 [22] _pullback(ctx::Zygote.Context{false}, f::typeof(f3), args::CuArray{Float32, 1, CUDA.Mem.DeviceBuffer})
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface2.jl:0
 [23] pullback(f::Function, cx::Zygote.Context{false}, args::CuArray{Float32, 1, CUDA.Mem.DeviceBuffer})
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface.jl:44
 [24] pullback
    @ C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface.jl:42 [inlined]
 [25] gradient(f::Function, args::CuArray{Float32, 1, CUDA.Mem.DeviceBuffer})
    @ Zygote C:\Users\Hossein Pourbozorg\.julia\packages\Zygote\YYT6v\src\compiler\interface.jl:96
 [26] top-level scope
    @ REPL[21]:1
 [27] top-level scope
    @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\nbRJk\src\initialization.jl:205

@prbzrg prbzrg mentioned this issue Nov 6, 2023
Closed
@maleadt maleadt added enhancement New feature or request help wanted Extra attention is needed good first issue Good for newcomers and removed bug Something isn't working help wanted Extra attention is needed labels Nov 6, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
enhancement New feature or request good first issue Good for newcomers
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
2 participants
@maleadt @prbzrg