Cleanup & remove some deps (#1560)

Author: pxl-th

Project.toml

@@ -10,7 +10,6 @@ DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
-GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 IRTools = "7869d1d1-7146-5819-86e3-90919afe41df"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
@@ -20,23 +19,25 @@ MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [weakdeps]
+Atom = "c52e3926-4ff0-5f6e-af25-54175e0327b1"
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 
 [extensions]
+ZygoteAtomExt = "Atom"
 ZygoteColorsExt = "Colors"
 ZygoteDistancesExt = "Distances"
 ZygoteTrackerExt = "Tracker"
 
 [compat]
+Atom = "0.12"
 AbstractFFTs = "1.3.1"
 ChainRules = "1.72.2"
 ChainRulesCore = "1.25.1"
@@ -45,14 +46,12 @@ DiffRules = "1.4"
 Distances = "0.10"
 FillArrays = "0.8, 0.9, 0.10, 0.11, 0.12, 0.13, 1"
 ForwardDiff = "0.10"
-GPUArrays = "8.4.2, 9, 10, 11"
 GPUArraysCore = "0.1.1, 0.2"
 IRTools = "0.4.12"
 LogExpFunctions = "0.3.1"
 MacroTools = "0.5"
 NaNMath = "0.3, 1"
 PrecompileTools = "1"
-Requires = "1.1"
 SpecialFunctions = "1.6, 2"
 Statistics = "1"
 Tracker = "0.2"

ext/ZygoteAtomExt.jl

@@ -0,0 +1,10 @@
+module ZygoteAtomExt
+
+using Atom
+using Zygote
+using Zygote.Profile
+
+Zygote.Profile.atom_expandpath(path::String) = Atom.expandpath(path)
+Zygote.Profile.juno(node::Zygote.Profile.Node) = Atom.msg("profile", Zygote.Profile.tojson(node))
+
+end

ext/ZygoteColorsExt.jl

@@ -1,12 +1,7 @@
 module ZygoteColorsExt
 
-if isdefined(Base, :get_extension)
-    using Zygote
-    using Colors
-else
-    using ..Zygote
-    using ..Colors
-end
+using Zygote
+using Colors
 
 Zygote.@non_differentiable Colors.ColorTypes._parameter_upper_bound(::Any...)
 

ext/ZygoteDistancesExt.jl

@@ -1,14 +1,8 @@
 module ZygoteDistancesExt
 
-if isdefined(Base, :get_extension)
-    using Zygote
-    using Distances
-    using LinearAlgebra
-else
-    using ..Zygote
-    using ..Distances
-    using ..LinearAlgebra
-end
+using Zygote
+using Distances
+using LinearAlgebra
 
 using Zygote: @adjoint, AContext, _pullback
 

ext/ZygoteTrackerExt.jl

@@ -1,12 +1,7 @@
 module ZygoteTrackerExt
 
-if isdefined(Base, :get_extension)
-    using Zygote
-    using Tracker: Tracker, TrackedArray, TrackedReal
-else
-    using ..Zygote
-    using ..Tracker: Tracker, TrackedArray, TrackedReal
-end
+using Zygote
+using Tracker: Tracker, TrackedArray, TrackedReal
 
 Zygote.unwrap(x::Union{TrackedArray,TrackedReal}) = Tracker.data(x)
 

src/Zygote.jl

@@ -10,7 +10,7 @@ import ZygoteRules: @adjoint, @adjoint!, AContext, adjoint, _pullback, pullback,
 using ChainRulesCore
 using ChainRules: ChainRules, AbstractThunk, rrule, unthunk, canonicalize
 using IRTools
-using MacroTools, Requires
+using MacroTools
 using MacroTools: @forward
 
 import Distributed: pmap, CachingPool, workers
@@ -53,22 +53,7 @@ include("compiler/interface2.jl")
 
 include("profiler/Profile.jl")
 
-
-if !isdefined(Base, :get_extension)
-  @init @require Distances="b4f34e82-e78d-54a5-968a-f98e89d6e8f7" include("../ext/ZygoteDistancesExt.jl")
-  @init @require Tracker="9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" include("../ext/ZygoteTrackerExt.jl")
-  @init @require Colors="5ae59095-9a9b-59fe-a467-6f913c188581" include("../ext/ZygoteColorsExt.jl")
-end
-
 using InteractiveUtils
-precompile() = Requires.@include("precompile.jl")
-
-# helps to work around 265-y issues
-function refresh()
-  Requires.@include("compiler/interface2.jl")
-  precompile()
-  return
-end
 
 macro profile(ex)
   @capture(ex, f_(x__)) || error("@profile f(args...)")
@@ -79,6 +64,8 @@ macro profile(ex)
 end
 
 using PrecompileTools
-  @compile_workload precompile()
+@compile_workload begin
+    include("precompile.jl")
+end
 
-end # module
+end

src/compiler/interface2.jl

@@ -62,8 +62,6 @@ function _generate_callable_pullback(j::Type{<:Pullback{T}}, world, Δ) where T
   return update!(meta.code, back)
 end
 
-if VERSION >= v"1.10.0-DEV.873"
-
 # on Julia 1.10, generated functions need to keep track of the world age
 
 function _pullback_generator(world::UInt, source, self, ctx, f, args)
@@ -103,15 +101,3 @@ end
   $(Expr(:meta, :generated, _callable_pullback_generator))
   $(Expr(:meta, :generated_only))
 end
-
-else
-
-@generated function _pullback(ctx::AContext, f, args...)
-  _generate_pullback(ctx, nothing, f, args...)
-end
-
-@generated function (j::Pullback)(Δ)
-  _generate_callable_pullback(j, nothing, Δ)
-end
-
-end

src/profiler/Profile.jl

@@ -1,6 +1,5 @@
 module Profile
 
-using Requires
 using ..Zygote: Pullback, meta, stacklines
 
 function loc(f)
@@ -81,17 +80,18 @@ function profile(x)
   Node(Symbol(""), "", -1, cs)
 end
 
-@init @require Atom="c52e3926-4ff0-5f6e-af25-54175e0327b1" begin
-  function tojson(n::Node)
-    name, path = Atom.expandpath(string(n.file))
+# Defined in the ZygoteAtomExt.
+function atom_expandpath end
+function juno end
+
+function tojson(n::Node)
+    name, path = atom_expandpath(string(n.file))
     Dict(:path => path,
          :location => name,
          :func => string(n.func),
          :line => n.line,
          :count => n.size,
          :children => map(tojson, n.children))
-  end
-  juno(n::Node) = Atom.msg("profile", tojson(n))
 end
 
 end

test/chainrules_tests.jl

@@ -405,13 +405,8 @@ end
 
     # To test the generic case, we need a struct within a struct.
     nested = Tangent{Base.RefValue{ComplexF64}}(; x=Tangent{ComplexF64}(; re=1, im=NoTangent()),)
-    if VERSION > v"1.7-"
-        @test @inferred(Zygote.z2d((; x=(; re=1)), Ref(3.0+im))) == nested
-        @test @inferred(Zygote.z2d((; x=(; re=nothing)), Ref(3.0+im))) === NoTangent()
-    else
-        @test Zygote.z2d((; x=(; re=1)), Ref(3.0+im)) == nested
-        @test Zygote.z2d((; x=(; re=nothing)), Ref(3.0+im)) === NoTangent()
-    end
+    @test @inferred(Zygote.z2d((; x=(; re=1)), Ref(3.0+im))) == nested
+    @test @inferred(Zygote.z2d((; x=(; re=nothing)), Ref(3.0+im))) === NoTangent()
 
     x = (c = (a = randn(3,3), b = rand(3)), d = randn(5))
     z2d_compiled = Zygote.z2d(x, x)

test/compiler_tests.jl

@@ -40,8 +40,6 @@ bt = try back(1) catch e stacktrace(catch_backtrace()) end
 
 # Type inference checks
 
-Zygote.refresh()
-
 y, back = @test_inferred pullback(*, 2, 3)
 @test_inferred(back(1))