removes @pure and some miscelanious cleanups

Author: oscardssmith

Manifest.toml

@@ -27,9 +27,9 @@ uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 
 [[deps.ChainRules]]
 deps = ["Adapt", "ChainRulesCore", "Compat", "Distributed", "GPUArraysCore", "IrrationalConstants", "LinearAlgebra", "Random", "RealDot", "SparseArrays", "Statistics", "StructArrays"]
-git-tree-sha1 = "c46adabdd0348f0ee8de91142cfc4a72a613ac0a"
+git-tree-sha1 = "fdde4d8a31cf82b1d136cf6cb53924e8744a832b"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "1.46.1"
+version = "1.47.0"
 
 [[deps.ChainRulesCore]]
 deps = ["Compat", "LinearAlgebra", "SparseArrays"]
@@ -265,9 +265,9 @@ version = "1.10.0"
 
 [[deps.StaticArrays]]
 deps = ["LinearAlgebra", "Random", "StaticArraysCore", "Statistics"]
-git-tree-sha1 = "129703d62117c374c4f2db6d13a027741c46eafd"
+git-tree-sha1 = "cee507162ecbb677450f20058ca83bd559b6b752"
 uuid = "90137ffa-7385-5640-81b9-e52037218182"
-version = "1.5.13"
+version = "1.5.14"
 
 [[deps.StaticArraysCore]]
 git-tree-sha1 = "6b7ba252635a5eff6a0b0664a41ee140a1c9e72a"

src/extra_rules.jl

@@ -79,7 +79,7 @@ function (::∂⃖{N})(f::typeof(*), args...) where {N}
         end
         return z
     else
-        ∂⃖p = ∂⃖{minus1(N)}()
+        ∂⃖p = ∂⃖{N-1}()
         @destruct z, z̄ = ∂⃖p(rrule_times, f, args...)
         if z === nothing
             return ∂⃖recurse{N}()(f, args...)
@@ -130,15 +130,15 @@ end
 struct NonDiffEven{N, O, P}; end
 struct NonDiffOdd{N, O, P}; end
 
-(::NonDiffOdd{N, O, P})(Δ) where {N, O, P} = (ntuple(_->ZeroTangent(), N), NonDiffEven{N, plus1(O), P}())
-(::NonDiffEven{N, O, P})(Δ...) where {N, O, P} = (ZeroTangent(), NonDiffOdd{N, plus1(O), P}())
+(::NonDiffOdd{N, O, P})(Δ) where {N, O, P} = (ntuple(_->ZeroTangent(), N), NonDiffEven{N, O+1, P}())
+(::NonDiffEven{N, O, P})(Δ...) where {N, O, P} = (ZeroTangent(), NonDiffOdd{N, O+1, P}())
 (::NonDiffOdd{N, O, O})(Δ) where {N, O} = ntuple(_->ZeroTangent(), N)
 
 # This should not happen
 (::NonDiffEven{N, O, O})(Δ...) where {N, O} = error()
 
-@Base.pure function ChainRulesCore.rrule(::DiffractorRuleConfig, ::typeof(Core.apply_type), head, args...)
-    Core.apply_type(head, args...), NonDiffOdd{plus1(plus1(length(args))), 1, 1}()
+@Base.assume_effects :total function ChainRulesCore.rrule(::DiffractorRuleConfig, ::typeof(Core.apply_type), head, args...)
+    Core.apply_type(head, args...), NonDiffOdd{length(args)+2, 1, 1}()
 end
 
 function ChainRulesCore.rrule(::DiffractorRuleConfig, ::typeof(Core.tuple), args...)

src/interface.jl

@@ -143,11 +143,9 @@ Base.show(io::IO, f::PrimeDerivativeBack{N}) where {N} = print(io, f.f, "'"^N)
 
 # This improves performance for nested derivatives by short cutting some
 # recursion into the PrimeDerivative constructor
-@Base.pure minus1(N) = N - 1
-@Base.pure plus1(N) = N + 1
-lower_pd(f::PrimeDerivativeBack{N,T}) where {N,T} = PrimeDerivativeBack{minus1(N),T}(getfield(f, :f))
+lower_pd(f::PrimeDerivativeBack{N,T}) where {N,T} = PrimeDerivativeBack{N-1,T}(getfield(f, :f))
 lower_pd(f::PrimeDerivativeBack{1}) = getfield(f, :f)
-raise_pd(f::PrimeDerivativeBack{N,T}) where {N,T} = PrimeDerivativeBack{plus1(N),T}(getfield(f, :f))
+raise_pd(f::PrimeDerivativeBack{N,T}) where {N,T} = PrimeDerivativeBack{N+1,T}(getfield(f, :f))
 
 ChainRulesCore.rrule(::typeof(lower_pd), f) = lower_pd(f), Δ->(ZeroTangent(), Δ)
 ChainRulesCore.rrule(::typeof(raise_pd), f) = raise_pd(f), Δ->(ZeroTangent(), Δ)
@@ -170,8 +168,8 @@ end
 PrimeDerivativeFwd(f) = PrimeDerivativeFwd{1, typeof(f)}(f)
 PrimeDerivativeFwd(f::PrimeDerivativeFwd{N, T}) where {N, T} = raise_pd(f)
 
-lower_pd(f::PrimeDerivativeFwd{N,T}) where {N,T} = (error(); PrimeDerivativeFwd{minus1(N),T}(getfield(f, :f)))
-raise_pd(f::PrimeDerivativeFwd{N,T}) where {N,T} = PrimeDerivativeFwd{plus1(N),T}(getfield(f, :f))
+lower_pd(f::PrimeDerivativeFwd{N,T}) where {N,T} = (error(); PrimeDerivativeFwd{N-1,T}(getfield(f, :f)))
+raise_pd(f::PrimeDerivativeFwd{N,T}) where {N,T} = PrimeDerivativeFwd{N+1,T}(getfield(f, :f))
 
 (f::PrimeDerivativeFwd{0})(x) = getfield(f, :f)(x)
 

src/stage1/forward.jl

@@ -23,13 +23,13 @@ my_frule(::ZeroBundle{1, typeof(my_frule)}, args::ATB{1}...) = nothing
 (::∂☆{N})(::ZeroBundle{N, typeof(my_frule)}, ::ZeroBundle{N, ZeroBundle{1, typeof(my_frule)}}, args::ATB{N}...) where {N} = ZeroBundle{N}(nothing)
 
 shuffle_down(b::UniformBundle{N, B, U}) where {N, B, U} =
-    UniformBundle{minus1(N), <:Any, U}(UniformBundle{1, B, U}(b.primal, b.tangent.val), b.tangent.val)
+    UniformBundle{N-1, <:Any, U}(UniformBundle{1, B, U}(b.primal, b.tangent.val), b.tangent.val)
 
 function shuffle_down(b::ExplicitTangentBundle{N, B}) where {N, B}
     # N.B: This depends on the special properties of the canonical tangent index order
     ExplicitTangentBundle{N-1}(
         ExplicitTangentBundle{1}(b.primal, (partial(b, 1),)),
-        ntuple(2^(N-1)-1) do i
+        ntuple(1<<(N-1)-1) do i
             ExplicitTangentBundle{1}(partial(b, 2*i), (partial(b, 2*i+1),))
         end)
 end
@@ -86,7 +86,7 @@ function shuffle_up(r::CompositeBundle{N}) where {N}
     else
         return TangentBundle{N+1}(r.tup[1].primal,
             (r.tup[1].tangent.partials..., primal(b),
-            ntuple(i->partial(b,i), 2^(N+1)-1)...))
+            ntuple(i->partial(b,i), 1<<(N+1)-1)...))
     end
 end
 
@@ -124,8 +124,8 @@ function ChainRulesCore.frule_via_ad(::DiffractorRuleConfig, partials, args...)
 end
 
 function (::∂☆shuffle{N})(args::AbstractTangentBundle{N}...) where {N}
-    ∂☆p = ∂☆{minus1(N)}()
-    ∂☆p(ZeroBundle{minus1(N)}(my_frule), map(shuffle_down, args)...)
+    ∂☆p = ∂☆{N-1}()
+    ∂☆p(ZeroBundle{N-1}(my_frule), map(shuffle_down, args)...)
 end
 
 function (::∂☆internal{N})(args::AbstractTangentBundle{N}...) where {N}

src/stage1/generated.jl

@@ -101,12 +101,12 @@ struct ∂⃖weaveInnerOdd{N, O}; b̄; end
 end
 @Base.constprop :aggressive function (w::∂⃖weaveInnerOdd{N, O})(Δ) where {N, O}
     @destruct c, c̄ = w.b̄(Δ...)
-    return (c̄, c), ∂⃖weaveInnerEven{plus1(N), O}()
+    return (c̄, c), ∂⃖weaveInnerEven{N+1, O}()
 end
 struct ∂⃖weaveInnerEven{N, O}; end
 @Base.constprop :aggressive function (w::∂⃖weaveInnerEven{N, O})(Δ′, x...) where {N, O}
     @destruct y, ȳ  = Δ′(x...)
-    return y, ∂⃖weaveInnerOdd{plus1(N), O}(ȳ)
+    return y, ∂⃖weaveInnerOdd{N+1, O}(ȳ)
 end
 
 struct ∂⃖weaveOuterOdd{N, O}; end
@@ -115,15 +115,15 @@ struct ∂⃖weaveOuterOdd{N, O}; end
 end
 @Base.constprop :aggressive function (w::∂⃖weaveOuterOdd{N, O})((Δ′′, Δ′′′)) where {N, O}
     @destruct α, ᾱ = Δ′′′(Δ′′)
-    return (NoTangent(), α...), ∂⃖weaveOuterEven{plus1(N), O}(ᾱ)
+    return (NoTangent(), α...), ∂⃖weaveOuterEven{N+1, O}(ᾱ)
 end
 struct ∂⃖weaveOuterEven{N, O}; ᾱ end
 @Base.constprop :aggressive function (w::∂⃖weaveOuterEven{N, O})(Δ⁴...) where {N, O}
-    return w.ᾱ(Base.tail(Δ⁴)...), ∂⃖weaveOuterOdd{plus1(N), O}()
+    return w.ᾱ(Base.tail(Δ⁴)...), ∂⃖weaveOuterOdd{N+1, O}()
 end
 
 function (::∂⃖{N})(::∂⃖{1}, args...) where {N}
-    @destruct (a, ā) = ∂⃖{plus1(N)}()(args...)
+    @destruct (a, ā) = ∂⃖{N+1}()(args...)
     let O = c_order(N)
         (a, Protected{N}(@opaque Δ->begin
                 (b, b̄) = ā(Δ)
@@ -188,10 +188,10 @@ end
 (::∂⃖rruleD{N, N})(Δ...) where {N} = error("Should not be reached")
 
 # ∂⃖rrule
-@Base.pure term_depth(N) = 2^(N-2)
+term_depth(N) = 1<<(N-2)
 function (::∂⃖rrule{N})(z, z̄) where {N}
     @destruct (y, ȳ) = z
-    y, ∂⃖rruleA{term_depth(N), 1}(∂⃖{minus1(N)}(), ȳ, z̄)
+    y, ∂⃖rruleA{term_depth(N), 1}(∂⃖{N-1}(), ȳ, z̄)
 end
 
 function (::∂⃖{N})(f::Core.IntrinsicFunction, args...) where {N}
@@ -217,7 +217,7 @@ function (::∂⃖{N})(f::T, args...) where {T, N}
         end
         return z
     else
-        ∂⃖p = ∂⃖{minus1(N)}()
+        ∂⃖p = ∂⃖{N-1}()
         @destruct z, z̄ = ∂⃖p(rrule, f, args...)
         if z === nothing
             return ∂⃖recurse{N}()(f, args...)
@@ -231,7 +231,7 @@ function ChainRulesCore.rrule_via_ad(::DiffractorRuleConfig, f::T, args...) wher
     Tuple{Any, Any}(∂⃖{1}()(f, args...))
 end
 
-@Base.pure function (::∂⃖{1})(::typeof(Core.apply_type), head, args...)
+@Base.assume_effects :total function (::∂⃖{1})(::typeof(Core.apply_type), head, args...)
     return rrule(Core.apply_type, head, args...)
 end
 
@@ -284,8 +284,8 @@ struct EvenOddEven{O, P, F, G}; f::F; g::G; end
 EvenOddEven{O, P}(f::F, g::G) where {O, P, F, G} = EvenOddEven{O, P, F, G}(f, g)
 struct EvenOddOdd{O, P, F, G}; f::F; g::G; end
 EvenOddOdd{O, P}(f::F, g::G) where {O, P, F, G} = EvenOddOdd{O, P, F, G}(f, g)
-@Base.constprop :aggressive (o::EvenOddOdd{O, P, F, G})(Δ) where {O, P, F, G} = (o.f(Δ), EvenOddEven{plus1(O), P, F, G}(o.f, o.g))
-@Base.constprop :aggressive (e::EvenOddEven{O, P, F, G})(Δ...) where {O, P, F, G} = (e.g(Δ...), EvenOddOdd{plus1(O), P, F, G}(e.f, e.g))
+@Base.constprop :aggressive (o::EvenOddOdd{O, P, F, G})(Δ) where {O, P, F, G} = (o.f(Δ), EvenOddEven{O+1, P, F, G}(o.f, o.g))
+@Base.constprop :aggressive (e::EvenOddEven{O, P, F, G})(Δ...) where {O, P, F, G} = (e.g(Δ...), EvenOddOdd{O+1, P, F, G}(e.f, e.g))
 @Base.constprop :aggressive (o::EvenOddOdd{O, O})(Δ) where {O} = o.f(Δ)
 
 
@@ -363,11 +363,11 @@ struct ApplyOdd{O, P}; u; ∂⃖f; end
 struct ApplyEven{O, P}; u; ∂⃖∂⃖f; end
 @Base.constprop :aggressive function (a::ApplyOdd{O, P})(Δ) where {O, P}
     r, ∂⃖∂⃖f = a.∂⃖f(Δ)
-    (a.u(r), ApplyEven{plus1(O), P}(a.u, ∂⃖∂⃖f))
+    (a.u(r), ApplyEven{O+1, P}(a.u, ∂⃖∂⃖f))
 end
 @Base.constprop :aggressive function (a::ApplyEven{O, P})(_, _, ff, args...) where {O, P}
     r, ∂⃖∂⃖∂⃖f = Core._apply_iterate(iterate, a.∂⃖∂⃖f, (ff,), args...)
-    (r, ApplyOdd{plus1(O), P}(a.u, ∂⃖∂⃖∂⃖f))
+    (r, ApplyOdd{O+1, P}(a.u, ∂⃖∂⃖∂⃖f))
 end
 @Base.constprop :aggressive function (a::ApplyOdd{O, O})(Δ) where {O}
     r = a.∂⃖f(Δ)
@@ -381,10 +381,10 @@ function (this::∂⃖{N})(::typeof(Core._apply_iterate), iterate, f, args::Unio
 end
 
 
-@Base.pure c_order(N::Int) = 2^N - 1
+c_order(N::Int) = 1<<N - 1
 
-@Base.pure function (::∂⃖{N})(::typeof(Core.apply_type), head, args...) where {N}
-    Core.apply_type(head, args...), NonDiffOdd{plus1(plus1(length(args))), 1, c_order(N)}()
+@Base.assume_effects :total function (::∂⃖{N})(::typeof(Core.apply_type), head, args...) where {N}
+    Core.apply_type(head, args...), NonDiffOdd{length(args)+2, 1, c_order(N)}()
 end
 
 @Base.constprop :aggressive lifted_getfield(x, s) = getfield(x, s)

src/stage1/mixed.jl

@@ -14,13 +14,13 @@ end
 function (x::∂⃖composeOdd)(Δ)
     b, ∂b = x.b(Δ)
     a, ∂a = x.a(b[end])
-    a, ∂⃖composeEven{N, plus1(N)}(∂a, ∂b)
+    a, ∂⃖composeEven{N, N+1}(∂a, ∂b)
 end
 
 function (x::∂⃖composeEven)(args...)
     a, ∂a = x.a(args...)
     b, ∂b = x.b(a)
-    b, ∂⃖composeOdd{N, plus1(N)}(∂a, ∂b)
+    b, ∂⃖composeOdd{N, N+1}(∂a, ∂b)
 end
 
 function (x::∂⃖composeOdd{N,N})(Δ) where {N}