Use JuliaFormatter

.JuliaFormatter.toml

@@ -0,0 +1 @@
+style = "blue"

.github/workflows/CI.yml

@@ -21,7 +21,7 @@ jobs:
       fail-fast: false
       matrix:
         test_group: [
-          'aqua',
+          'quality',
           'basic',
           'rrules/avoiding_non_differentiable_code',
           'rrules/blas',

Project.toml

@@ -1,7 +1,7 @@
 name = "Mooncake"
 uuid = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 authors = ["Will Tebbutt, Hong Ge, and contributors"]
-version = "0.4.49"
+version = "0.4.50"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -55,6 +55,7 @@ FunctionWrappers = "1.1.3"
 Graphs = "1"
 InteractiveUtils = "1"
 JET = "0.9"
+JuliaFormatter = "1.0"
 LinearAlgebra = "1"
 LuxLib = "1"
 MistyClosures = "2"
@@ -74,9 +75,10 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 DiffTests = "de460e47-3fe3-5279-bb4a-814414816d5d"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
+JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["AllocCheck", "Aqua", "BenchmarkTools", "DiffTests", "JET", "Pkg", "StableRNGs", "Test"]
+test = ["AllocCheck", "Aqua", "BenchmarkTools", "DiffTests", "JET", "JuliaFormatter", "Pkg", "StableRNGs", "Test"]

bench/run_benchmarks.jl

@@ -1,8 +1,7 @@
 using Pkg
-Pkg.develop(path=joinpath(@__DIR__, ".."))
+Pkg.develop(; path=joinpath(@__DIR__, ".."))
 
-using
-    AbstractGPs,
+using AbstractGPs,
     Chairmarks,
     CSV,
     DataFrames,
@@ -28,13 +27,13 @@ using Mooncake:
 
 using Mooncake.TestUtils: _deepcopy
 
-function to_benchmark(__rrule!!::R, dx::Vararg{CoDual, N}) where {R, N}
+function to_benchmark(__rrule!!::R, dx::Vararg{CoDual,N}) where {R,N}
     dx_f = Mooncake.tuple_map(x -> CoDual(primal(x), Mooncake.fdata(tangent(x))), dx)
     out, pb!! = __rrule!!(dx_f...)
     return pb!!(Mooncake.zero_rdata(primal(out)))
 end
 
-function zygote_to_benchmark(ctx, x::Vararg{Any, N}) where {N}
+function zygote_to_benchmark(ctx, x::Vararg{Any,N}) where {N}
     out, pb = Zygote._pullback(ctx, x...)
     return pb(out)
 end
@@ -107,7 +106,7 @@ end
 @model broadcast_demo(x) = begin
     μ ~ truncated(Normal(1, 2), 0.1, 10)
     σ ~ truncated(Normal(1, 2), 0.1, 10)
-    x .~ LogNormal(μ, σ)   
+    x .~ LogNormal(μ, σ)
 end
 
 function build_turing_problem()
@@ -122,17 +121,21 @@ function build_turing_problem()
     return test_function, randn(rng, d)
 end
 
-run_turing_problem(f::F, x::X) where {F, X} = f(x)
+run_turing_problem(f::F, x::X) where {F,X} = f(x)
 
-should_run_benchmark(
+function should_run_benchmark(
     ::Val{:zygote}, ::Base.Fix1{<:typeof(DynamicPPL.LogDensityProblems.logdensity)}, x...
-) = false
-should_run_benchmark(
+)
+    return false
+end
+function should_run_benchmark(
     ::Val{:enzyme}, ::Base.Fix1{<:typeof(DynamicPPL.LogDensityProblems.logdensity)}, x...
-) = false
+)
+    return false
+end
 should_run_benchmark(::Val{:enzyme}, x...) = false
 
-@inline g(x, a, ::Val{N}) where {N} = N > 0 ? g(x * a, a, Val(N-1)) : x
+@inline g(x, a, ::Val{N}) where {N} = N > 0 ? g(x * a, a, Val(N - 1)) : x
 
 large_single_block(x::AbstractVector{<:Real}) = g(x[1], x[2], Val(400))
 
@@ -168,14 +171,12 @@ function generate_inter_framework_tests()
 end
 
 function benchmark_rules!!(test_case_data, default_ratios, include_other_frameworks::Bool)
-
     test_cases = reduce(vcat, map(first, test_case_data))
     memory = map(x -> x[2], test_case_data)
     ranges = reduce(vcat, map(x -> x[3], test_case_data))
     tags = reduce(vcat, map(x -> x[4], test_case_data))
     GC.@preserve memory begin
         return map(enumerate(test_cases)) do (n, args)
-
             @info "$n / $(length(test_cases))", _typeof(args)
             suite = Dict()
 
@@ -186,7 +187,7 @@ function benchmark_rules!!(test_case_data, default_ratios, include_other_framewo
                 () -> primals,
                 primals -> (primals[1], _deepcopy(primals[2:end])),
                 (a -> a[1]((a[2]...))),
-                _ -> true,
+                _ -> true;
                 evals=1,
             )
 
@@ -199,17 +200,19 @@ function benchmark_rules!!(test_case_data, default_ratios, include_other_framewo
                 () -> (rule, coduals),
                 identity,
                 a -> to_benchmark(a[1], a[2]...),
-                _ -> true,
+                _ -> true;
                 evals=1,
             )
 
             if include_other_frameworks
-
                 if should_run_benchmark(Val(:zygote), args...)
                     @info "Zygote"
                     suite["zygote"] = @be(
-                        _, _, zygote_to_benchmark($(Zygote.Context()), $primals...), _,
-                        evals=1,
+                        _,
+                        _,
+                        zygote_to_benchmark($(Zygote.Context()), $primals...),
+                        _,
+                        evals = 1,
                     )
                 end
 
@@ -219,21 +222,27 @@ function benchmark_rules!!(test_case_data, default_ratios, include_other_framewo
                     compiled_tape = ReverseDiff.compile(tape)
                     result = map(x -> randn(size(x)), primals[2:end])
                     suite["rd"] = @be(
-                        _, _, rd_to_benchmark!($result, $compiled_tape, $primals[2:end]), _,
-                        evals=1,
+                        _,
+                        _,
+                        rd_to_benchmark!($result, $compiled_tape, $primals[2:end]),
+                        _,
+                        evals = 1,
                     )
                 end
 
                 if should_run_benchmark(Val(:enzyme), args...)
                     @info "Enzyme"
                     dup_args = map(x -> Duplicated(x, randn(size(x))), primals[2:end])
                     suite["enzyme"] = @be(
-                        _, _, autodiff(Reverse, $primals[1], Active, $dup_args...), _,
-                        evals=1,
+                        _,
+                        _,
+                        autodiff(Reverse, $primals[1], Active, $dup_args...),
+                        _,
+                        evals = 1,
                     )
                 end
             end
-            
+
             return combine_results((args, suite), tags[n], ranges[n], default_ratios)
         end
     end
@@ -319,7 +328,7 @@ well-suited to the numbers typically found in this field.
 function plot_ratio_histogram!(df::DataFrame)
     bin = 10.0 .^ (-1.0:0.05:4.0)
     xlim = extrema(bin)
-    histogram(df.Mooncake; xscale=:log10, xlim, bin, title="log", label="")
+    return histogram(df.Mooncake; xscale=:log10, xlim, bin, title="log", label="")
 end
 
 function create_inter_ad_benchmarks()
@@ -328,7 +337,7 @@ function create_inter_ad_benchmarks()
     df = DataFrame(results)[:, [:tag, tools...]]
 
     # Plot graph of results.
-    plt = plot(yscale=:log10, legend=:topright, title="AD Time / Primal Time (Log Scale)")
+    plt = plot(; yscale=:log10, legend=:topright, title="AD Time / Primal Time (Log Scale)")
     for label in string.(tools)
         plot!(plt, df.tag, df[:, label]; label, marker=:circle, xrotation=45)
     end
@@ -337,7 +346,9 @@ function create_inter_ad_benchmarks()
     # Write table of results.
     formatted_cols = map(t -> t => string.(round.(df[:, t]; sigdigits=3)), tools)
     df_formatted = DataFrame(:Label => df.tag, formatted_cols...)
-    open(io -> pretty_table(io, df_formatted), "bench/benchmark_results.txt"; write=true)
+    return open(
+        io -> pretty_table(io, df_formatted), "bench/benchmark_results.txt"; write=true
+    )
 end
 
 function main()

docs/make.jl

@@ -9,26 +9,18 @@ DocMeta.setdocmeta!(
     recursive=true,
 )
 
-makedocs(
+makedocs(;
     sitename="Mooncake.jl",
     format=Documenter.HTML(;
-        mathengine = Documenter.KaTeX(
-            Dict(
-                :macros => Dict(
-                    "\\RR" => "\\mathbb{R}",
-                ),
-            )
-        ),
+        mathengine=Documenter.KaTeX(Dict(:macros => Dict("\\RR" => "\\mathbb{R}"))),
         size_threshold_ignore=[
-            joinpath("developer_documentation", "internal_docstrings.md"),
+            joinpath("developer_documentation", "internal_docstrings.md")
         ],
     ),
     modules=[Mooncake],
     checkdocs=:none,
-    plugins=[
-        CitationBibliography(joinpath(@__DIR__, "src", "refs.bib"); style=:numeric),
-    ],
-    pages = [
+    plugins=[CitationBibliography(joinpath(@__DIR__, "src", "refs.bib"); style=:numeric)],
+    pages=[
         "Mooncake.jl" => "index.md",
         "Understanding Mooncake.jl" => [
             joinpath("understanding_mooncake", "introduction.md"),
@@ -46,7 +38,7 @@ makedocs(
             joinpath("developer_documentation", "internal_docstrings.md"),
         ],
         "known_limitations.md",
-    ]
+    ],
 )
 
-deploydocs(repo="github.com/compintell/Mooncake.jl.git", push_preview=true)
+deploydocs(; repo="github.com/compintell/Mooncake.jl.git", push_preview=true)

ext/MooncakeAllocCheckExt.jl

@@ -3,6 +3,6 @@ module MooncakeAllocCheckExt
 using AllocCheck, Mooncake
 import Mooncake.TestUtils: check_allocs, Shim
 
-@check_allocs check_allocs(::Shim, f::F, x::Tuple{Vararg{Any, N}}) where {F, N} = f(x...)
+@check_allocs check_allocs(::Shim, f::F, x::Tuple{Vararg{Any,N}}) where {F,N} = f(x...)
 
 end

ext/MooncakeCUDAExt.jl

@@ -38,7 +38,7 @@
 _add_to_primal(x::P, y::P, ::Bool) where {P<:CuArray{<:IEEEFloat}} = x + y
 _diff(x::P, y::P) where {P<:CuArray{<:IEEEFloat}} = x - y
 _dot(x::P, y::P) where {P<:CuArray{<:IEEEFloat}} = Float64(dot(x, y))
-_scale(x::Float64, y::P) where {T<:IEEEFloat, P<:CuArray{T}} = T(x) * y
+_scale(x::Float64, y::P) where {T<:IEEEFloat,P<:CuArray{T}} = T(x) * y
 function populate_address_map!(m::AddressMap, p::CuArray, t::CuArray)
     k = pointer_from_objref(p)
     v = pointer_from_objref(t)
@@ -55,9 +55,7 @@
 
 # Basic rules for operating on CuArrays.
 
-@is_primitive(
-    MinimalCtx, Tuple{Type{<:CuArray}, UndefInitializer, Vararg{Int, N}} where {N},
-)
+@is_primitive(MinimalCtx, Tuple{Type{<:CuArray},UndefInitializer,Vararg{Int,N}} where {N},)
 function rrule!!(
     p::CoDual{Type{P}}, init::CoDual{UndefInitializer}, dims::CoDual{Int}...
 ) where {P<:CuArray{<:Base.IEEEFloat}}

ext/MooncakeDynamicPPLExt.jl

@@ -4,6 +4,6 @@ using DynamicPPL: DynamicPPL, istrans
 using Mooncake: Mooncake
 
 # This is purely an optimisation.
-Mooncake.@zero_adjoint Mooncake.DefaultCtx Tuple{typeof(istrans), Vararg}
+Mooncake.@zero_adjoint Mooncake.DefaultCtx Tuple{typeof(istrans),Vararg}
 
 end # module