Run Aqua Checks (#327)

* Add Aqua and sort out compat * Resolve method ambiguities * Remove type piracy * Bump patch * Add Aqua badge to readme
compintell · Oct 30, 2024 · b30d313 · b30d313 · willtebbutt · Oct 30, 2024
1 parent 1c91d69
commit b30d313
Show file tree

Hide file tree

Showing 9 changed files with 51 additions and 29 deletions.
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -21,6 +21,7 @@ jobs:
       fail-fast: false
       matrix:
         test_group:
+          - 'aqua'
           - 'basic'
           - 'rrules'
           - 'ext/differentiation_interface'

diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "Mooncake"
 uuid = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 authors = ["Will Tebbutt, Hong Ge, and contributors"]
-version = "0.4.25"
+version = "0.4.26"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -41,6 +41,7 @@ MooncakeSpecialFunctionsExt = "SpecialFunctions"
 [compat]
 ADTypes = "1.9"
 AllocCheck = "0.2"
+Aqua = "0.8.9"
 BenchmarkTools = "1"
 CUDA = "5"
 ChainRulesCore = "1"
@@ -49,19 +50,24 @@ DiffTests = "0.1"
 DynamicPPL = "0.29, 0.30"
 ExprTools = "0.1"
 Graphs = "1"
+InteractiveUtils = "1"
 JET = "0.9"
+LinearAlgebra = "1"
 LogDensityProblemsAD = "1"
 LuxLib = "1.2 - 1.3.3"
 MistyClosures = "2"
 NNlib = "0.9"
+Random = "1"
 Pkg = "1"
 Setfield = "1"
 SpecialFunctions = "2"
 StableRNGs = "1"
+Test = "1"
 julia = "1"
 
 [extras]
 AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a"
+Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 DiffTests = "de460e47-3fe3-5279-bb4a-814414816d5d"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
@@ -70,4 +76,4 @@ StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["AllocCheck", "BenchmarkTools", "DiffTests", "JET", "Pkg", "StableRNGs", "Test"]
+test = ["AllocCheck", "Aqua", "BenchmarkTools", "DiffTests", "JET", "Pkg", "StableRNGs", "Test"]
diff --git a/README.md b/README.md
@@ -5,6 +5,7 @@
 [![Code Style: Blue](https://img.shields.io/badge/code%20style-blue-4495d1.svg)](https://github.com/invenia/BlueStyle)
 [![ColPrac: Contributor's Guide on Collaborative Practices for Community Packages](https://img.shields.io/badge/ColPrac-Contributor's%20Guide-blueviolet)](https://github.com/SciML/ColPrac)
 [![](https://img.shields.io/badge/docs-blue.svg)](https://compintell.github.io/Mooncake.jl/dev)
+[![Aqua QA](https://raw.githubusercontent.com/JuliaTesting/Aqua.jl/master/badge.svg)](https://github.com/JuliaTesting/Aqua.jl)
 
 The goal of the `Mooncake.jl` project is to produce a reverse-mode AD package which is written entirely in Julia, which improves over both `ReverseDiff.jl` and `Zygote.jl` in several ways, and is competitive with `Enzyme.jl`.
 

diff --git a/src/interpreter/ir_utils.jl b/src/interpreter/ir_utils.jl
@@ -179,8 +179,9 @@ function optimise_ir!(ir::IRCode; show_ir=false, do_inline=true)
     return ir
 end
 
-Base.iterate(x::CC.MethodLookupResult) = CC.iterate(x)
-Base.iterate(x::CC.MethodLookupResult, n::Int) = CC.iterate(x, n)
+# Handles difference between 1.10 and 1.11.
+get_matches(x::CC.MethodLookupResult) = x.matches
+get_matches(x::Vector{Any}) = x
 
 """
     lookup_ir(
@@ -196,7 +197,7 @@ Returns a tuple containing the `IRCode` and its return type.
 function lookup_ir(interp::CC.AbstractInterpreter, tt::Type{<:Tuple}; optimize_until=nothing)
     matches = CC.findall(tt, CC.method_table(interp))
     asts = []
-    for match in matches.matches
+    for match in get_matches(matches.matches)
         match = match::Core.MethodMatch
         if VERSION < v"1.11-"
             meth = Base.func_for_method_checked(match.method, tt, match.sparams)

diff --git a/src/rrules/memory.jl b/src/rrules/memory.jl
@@ -408,7 +408,7 @@ end
 
 function rrule!!(
     ::CoDual{typeof(lgetfield)},
-    x::CoDual{<:Memory},
+    x::CoDual{<:Memory, <:Memory},
     ::CoDual{Val{name}},
     ::CoDual{Val{order}},
 ) where {name, order}
@@ -420,7 +420,7 @@ end
 
 function rrule!!(
     ::CoDual{typeof(lgetfield)},
-    x::CoDual{<:MemoryRef},
+    x::CoDual{<:MemoryRef, <:MemoryRef},
     ::CoDual{Val{name}},
     ::CoDual{Val{order}},
 ) where {name, order}
@@ -432,7 +432,7 @@ end
 
 function rrule!!(
     ::CoDual{typeof(lgetfield)},
-    x::CoDual{<:Array},
+    x::CoDual{<:Array, <:Array},
     ::CoDual{Val{name}},
     ::CoDual{Val{order}},
 ) where {name, order}
@@ -444,14 +444,16 @@ end
 
 const _MemTypes = Union{Memory, MemoryRef, Array}
 
-function rrule!!(f::CoDual{typeof(lgetfield)}, x::CoDual{<:_MemTypes}, name::CoDual{<:Val})
+function rrule!!(
+    f::CoDual{typeof(lgetfield)}, x::CoDual{<:_MemTypes, <:_MemTypes}, name::CoDual{<:Val}
+)
     y, adj = rrule!!(f, x, name, zero_fcodual(Val(:not_atomic)))
     ternary_lgetfield_adjoint(dy) = adj(dy)[1:3]
     return y, ternary_lgetfield_adjoint
 end
 
 function rrule!!(
-    ::CoDual{typeof(getfield)}, x::CoDual{<:_MemTypes},
+    ::CoDual{typeof(getfield)}, x::CoDual{<:_MemTypes, <:_MemTypes},
     name::CoDual{<:Union{Int, Symbol}},
     order::CoDual{Symbol},
 )
@@ -466,15 +468,20 @@ function rrule!!(
 end
 
 function rrule!!(
-    f::CoDual{typeof(getfield)}, x::CoDual{<:_MemTypes}, name::CoDual{<:Union{Int, Symbol}}
+    f::CoDual{typeof(getfield)},
+    x::CoDual{<:_MemTypes, <:_MemTypes},
+    name::CoDual{<:Union{Int, Symbol}},
 )
     y, adj = rrule!!(f, x, name, zero_fcodual(:not_atomic))
     ternary_getfield_adjoint(dy) = adj(dy)[1:3]
     return y, ternary_getfield_adjoint
 end
 
 @inline function rrule!!(
-    ::CoDual{typeof(lsetfield!)}, value::CoDual{<:Array}, ::CoDual{Val{name}}, x::CoDual
+    ::CoDual{typeof(lsetfield!)},
+    value::CoDual{<:Array, <:Array},
+    ::CoDual{Val{name}},
+    x::CoDual,
 ) where {name}
     old_x = getfield(value.x, name)
     old_dx = getfield(value.dx, name)

diff --git a/src/rrules/misc.jl b/src/rrules/misc.jl
@@ -31,6 +31,9 @@
 @zero_adjoint MinimalCtx Tuple{typeof(Base.padding), DataType, Int}
 @zero_adjoint MinimalCtx Tuple{Type, TypeVar, Type}
 
+# Required to avoid an ambiguity.
+@zero_adjoint MinimalCtx Tuple{Type{Symbol}, TypeVar, Type}
+
 if VERSION >= v"1.11-"
     @zero_adjoint MinimalCtx Tuple{typeof(Random.hash_seed), Vararg}
     @zero_adjoint MinimalCtx Tuple{typeof(Base.dataids), Memory}
@@ -75,10 +78,10 @@ lgetfield(x, ::Val{f}) where {f} = getfield(x, f)
     return y, pb!!
 end
 
-@inline _get_fdata_field(_, t::Union{Tuple, NamedTuple}, f...) = getfield(t, f...)
-@inline _get_fdata_field(_, data::FData, f...) = val(getfield(data.data, f...))
-@inline _get_fdata_field(primal, ::NoFData, f...) = uninit_fdata(getfield(primal, f...))
-@inline _get_fdata_field(_, t::MutableTangent, f...) = fdata(val(getfield(t.fields, f...)))
+@inline _get_fdata_field(_, t::Union{Tuple, NamedTuple}, f) = getfield(t, f)
+@inline _get_fdata_field(_, data::FData, f) = val(getfield(data.data, f))
+@inline _get_fdata_field(primal, ::NoFData, f) = uninit_fdata(getfield(primal, f))
+@inline _get_fdata_field(_, t::MutableTangent, f) = fdata(val(getfield(t.fields, f)))
 
 increment_field_rdata!(dx::MutableTangent, ::NoRData, ::Val) = dx
 increment_field_rdata!(dx::NoFData, ::NoRData, ::Val) = dx

diff --git a/src/rrules/tasks.jl b/src/rrules/tasks.jl
@@ -33,12 +33,12 @@ rdata_type(::Type{TaskTangent}) = NoRData
 
 tangent(t::TaskTangent, ::NoRData) = t
 
-@inline function _get_fdata_field(_, t::TaskTangent, f...)
-    f === (:rngState0, ) && return NoFData()
-    f === (:rngState1, ) && return NoFData()
-    f === (:rngState2, ) && return NoFData()
-    f === (:rngState3, ) && return NoFData()
-    f === (:rngState4, ) && return NoFData()
+@inline function _get_fdata_field(_, t::TaskTangent, f)
+    f === :rngState0 && return NoFData()
+    f === :rngState1 && return NoFData()
+    f === :rngState2 && return NoFData()
+    f === :rngState3 && return NoFData()
+    f === :rngState4 && return NoFData()
     throw(error("Unhandled field $f"))
 end
 
@@ -53,7 +53,9 @@ function get_tangent_field(t::TaskTangent, f)
     throw(error("Unhandled field $f"))
 end
 
-function rrule!!(::CoDual{typeof(lgetfield)}, x::CoDual{Task}, ::CoDual{Val{f}}) where {f}
+const TaskCoDual = CoDual{Task, TaskTangent}
+
+function rrule!!(::CoDual{typeof(lgetfield)}, x::TaskCoDual, ::CoDual{Val{f}}) where {f}
     dx = x.dx
     function mutable_lgetfield_pb!!(dy)
         increment_field_rdata!(dx, dy, Val{f}())
@@ -63,14 +65,12 @@ function rrule!!(::CoDual{typeof(lgetfield)}, x::CoDual{Task}, ::CoDual{Val{f}})
     return y, mutable_lgetfield_pb!!
 end
 
-function rrule!!(::CoDual{typeof(getfield)}, x::CoDual{Task}, f::CoDual)
+function rrule!!(::CoDual{typeof(getfield)}, x::TaskCoDual, f::CoDual)
     return rrule!!(zero_fcodual(lgetfield), x, zero_fcodual(Val(primal(f))))
 end
 
-function rrule!!(
-    ::CoDual{typeof(lsetfield!)}, task::CoDual{Task}, name::CoDual, value::CoDual
-)
-    return lsetfield_rrule(task, name, value)
+function rrule!!(::CoDual{typeof(lsetfield!)}, task::TaskCoDual, name::CoDual, val::CoDual)
+    return lsetfield_rrule(task, name, val)
 end
 
 set_tangent_field!(t::TaskTangent, f, ::NoTangent) = NoTangent()

diff --git a/test/front_matter.jl b/test/front_matter.jl
@@ -1,5 +1,6 @@
 using
     AllocCheck,
+    Aqua,
     BenchmarkTools,
     DiffRules,
     JET,

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -1,7 +1,9 @@
 include("front_matter.jl")
 
 @testset "Mooncake.jl" begin
-    if test_group == "basic"
+    if test_group == "aqua"
+        Aqua.test_all(Mooncake)
+    elseif test_group == "basic"
         include("utils.jl")
         include("tangents.jl")
         include("fwds_rvs_data.jl")