diff --git a/.github/copilot-instructions.md b/.github/copilot-instructions.md deleted file mode 100644 index 6111b7439c..0000000000 --- a/.github/copilot-instructions.md +++ /dev/null @@ -1,14 +0,0 @@ -This repo implements an automatic differentiation package for the Julia language. - -Please follow these guidelines when contributing: - -## Code Standards - -- Never commit changes to `Project.toml` to PRs. -- Format modified files using `JuliaFormatter`, eg, `using JuliaFormatter; format(".")` from the root of the repo. - -## Key Guidelines -- Follow Go best practices and idiomatic patterns -- Maintain existing code structure and organisation -- Write unit tests for new functionality. -- Document public APIs and complex logic. Suggest changes to the `docs/` folder when appropriate diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 9d4e266024..23eb3309e6 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -68,6 +68,27 @@ jobs: arch: ${{ matrix.arch }} include-all-prereleases: false - uses: julia-actions/cache@v2 + - name: Adjust JET version based on Julia version + run: | + julia_version=$(julia --version | grep -oE '[0-9]+\.[0-9]+') + echo "Detected Julia version: $julia_version" + + # Remove any existing JET constraint from [compat] section + sed -i '/^JET = "[0-9]\.[0-9]/d' Project.toml + + # Compare version using version sort + if printf '%s\n' "1.12" "$julia_version" | sort -V | head -n1 | grep -q "1.12"; then + echo "Julia version >= 1.12, using latest JET version (no constraint)" + else + echo "Julia version < 1.12, setting JET = \"0.9\"" + # Add JET = "0.9" constraint in the [compat] section + sed -i '/^\[compat\]/a JET = "0.9"' Project.toml + fi + + # Show the current JET constraint (if any) + echo "Current JET constraint in Project.toml:" + grep -A 20 "^\[compat\]" Project.toml | grep "JET = " || echo "No JET constraint found" + shell: bash - uses: julia-actions/julia-buildpkg@v1 - uses: julia-actions/julia-runtest@v1 env: @@ -95,6 +116,7 @@ jobs: {test_type: 'integration_testing', label: 'array'}, {test_type: 'integration_testing', label: 'bijectors'}, {test_type: 'integration_testing', label: 'diff_tests'}, + {test_type: 'integration_testing', label: 'diffeq'}, {test_type: 'integration_testing', label: 'dispatch_doctor'}, {test_type: 'integration_testing', label: 'distributions'}, {test_type: 'integration_testing', label: 'dynamicppl'}, @@ -119,6 +141,27 @@ jobs: arch: ${{ matrix.arch }} include-all-prereleases: false - uses: julia-actions/cache@v2 + - name: Adjust JET version based on Julia version + run: | + julia_version=$(julia --version | grep -oE '[0-9]+\.[0-9]+') + echo "Detected Julia version: $julia_version" + + # Remove any existing JET constraint from [compat] section + sed -i '/^JET = "[0-9]\.[0-9]/d' Project.toml + + # Compare version using version sort + if printf '%s\n' "1.12" "$julia_version" | sort -V | head -n1 | grep -q "1.12"; then + echo "Julia version >= 1.12, using latest JET version (no constraint)" + else + echo "Julia version < 1.12, setting JET = \"0.9\"" + # Add JET = "0.9" constraint in the [compat] section + sed -i '/^\[compat\]/a JET = "0.9"' Project.toml + fi + + # Show the current JET constraint (if any) + echo "Current JET constraint in Project.toml:" + grep -A 20 "^\[compat\]" Project.toml | grep "JET = " || echo "No JET constraint found" + shell: bash - uses: julia-actions/julia-buildpkg@v1 - run: | if [ ${{ matrix.test_group.test_type }} == 'ext' ]; then diff --git a/Project.toml b/Project.toml index 3837aef17d..85ab049b24 100644 --- 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.152" +version = "0.4.154" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" @@ -61,20 +61,20 @@ FunctionWrappers = "1.1.3" GPUArraysCore = "0.1, 0.2" Graphs = "1" InteractiveUtils = "1" -JET = "0.9" +JET = "0.9, 0.10" JuliaFormatter = "1.0, 2.1" JuliaInterpreter = "0.9" LinearAlgebra = "1" LuxLib = "1" -MistyClosures = "2" MLDataDevices = "1.10.0" +MistyClosures = "2" NNlib = "0.9" Pkg = "1" Random = "1" SLEEFPirates = "0.6.43" SpecialFunctions = "2" -Static = "1.1.1" StableRNGs = "1" +Static = "1.1.1" Test = "1" julia = "~1.10.8, 1.11.6" diff --git a/docs/src/developer_documentation/ir_representation.md b/docs/src/developer_documentation/ir_representation.md index d9d03c469a..367c20819d 100644 --- a/docs/src/developer_documentation/ir_representation.md +++ b/docs/src/developer_documentation/ir_representation.md @@ -1,4 +1,4 @@ -# IR Representation +# IR Representations and Code Transformations Mooncake.jl works by transforming Julia's SSA-form (static single assignment) Intermediate Representation (IR), so a good understanding of Julia's IR is needed to understand Mooncake. Furthermore, Mooncake holds Julia's IR in a different data structure than the one usually used when producing code for reverse-mode AD. @@ -555,6 +555,12 @@ Reverse-mode AD makes extensive use of such transformations, so `BBCode` is curr There are efforts such as [this PR](https://github.com/JuliaLang/julia/pull/45305) to augment `IRCode` with the capability to manipulate the CFG structure in a convenient manner. Ideally these efforts will succeed, then we can do away with `BBCode`. +### Comparison with Alternative Approaches + +It's worth noting that other automatic differentiation systems have taken different approaches to IR manipulation. For example, [Diffractor.jl](https://github.com/JuliaDiff/Diffractor.jl) uses an "Optics" approach for IR transformations. + +For readers interested in learning more about Julia's IR representation beyond what's covered here, the [Scientific Programming in Julia course materials](https://github.com/JuliaTeachingCTU/Scientific-Programming-in-Julia/blob/2023W/docs/src/lecture_09/ircode.md) provide additional valuable context. + ## Docstrings ```@autodocs; canonical=true diff --git a/docs/src/developer_documentation/reverse_mode_design.md b/docs/src/developer_documentation/reverse_mode_design.md index 8ef4ab190e..04979822ff 100644 --- a/docs/src/developer_documentation/reverse_mode_design.md +++ b/docs/src/developer_documentation/reverse_mode_design.md @@ -72,6 +72,8 @@ Do line by line transformation of the statements and then possibly refresh the C Examples of how line-by-line transformations can be done, are defined in [`Mooncake.make_ad_stmts!`](@ref). The `IRCode` nodes are not explicitly documented in or . Might need completion of official docs, but Mooncake docs in the meantime. +For additional information about `IRCode` and `BBCode` data structures and transformation examples, see [IR Representations and Code Transformations](@ref). + Inlining pass can prevent us from using high-level rules by inlining the function (e.g. unrolling a loop). The contexts in [`interpreter/contexts.jl`](https://github.com/chalk-lab/Mooncake.jl/blob/src/interpreter/contexts.jl) are `MinimalCtx` (necessary for AD to work) and `DefaultCtx` (ensure that we hit all of the rules). Distinction between rules is not well maintained in Mooncake at the moment. diff --git a/src/fwds_rvs_data.jl b/src/fwds_rvs_data.jl index 0ad4981e3e..2cf7bceee6 100644 --- a/src/fwds_rvs_data.jl +++ b/src/fwds_rvs_data.jl @@ -258,6 +258,10 @@ function fdata(t::T) where {T} return F(fdata(t.fields)) end +function fdata(::Type{T}) where {T} + error("$T is a type. Perhaps you meant fdata_type($T) or fdata(instance_of_tangent)?") +end + function fdata(t::T) where {T<:PossiblyUninitTangent} F = fdata_type(T) return is_init(t) ? F(fdata(val(t))) : F() @@ -530,6 +534,10 @@ function rdata(t::T) where {T} return R(rdata(t.fields)) end +function rdata(::Type{T}) where {T} + error("$T is a type. Perhaps you meant rdata_type($T) or rdata(instance_of_tangent)?") +end + function rdata(t::T) where {T<:PossiblyUninitTangent} R = rdata_type(T) return is_init(t) ? R(rdata(val(t))) : R() diff --git a/src/rrules/blas.jl b/src/rrules/blas.jl index f7555226d5..f7fd44a9e5 100644 --- a/src/rrules/blas.jl +++ b/src/rrules/blas.jl @@ -1109,6 +1109,13 @@ function invertible_blas_matrices(rng::AbstractRNG, P::Type{<:BlasFloat}, p::Int end end +function positive_definite_blas_matrices(rng::AbstractRNG, P::Type{<:BlasFloat}, p::Int) + return map(blas_matrices(rng, P, p, p)) do A + A .= A'A + I + return A + end +end + function blas_vectors(rng::AbstractRNG, P::Type{<:BlasFloat}, p::Int) xs = Any[ randn(rng, P, p), diff --git a/src/rrules/lapack.jl b/src/rrules/lapack.jl index ede01c2217..2d5229e21c 100644 --- a/src/rrules/lapack.jl +++ b/src/rrules/lapack.jl @@ -1,14 +1,14 @@ # See https://sethaxen.com/blog/2021/02/differentiating-the-lu-decomposition/ for details. -@is_primitive(MinimalCtx, Tuple{typeof(LAPACK.getrf!),AbstractMatrix{<:BlasRealFloat}}) +@is_primitive(MinimalCtx, Tuple{typeof(LAPACK.getrf!),AbstractMatrix{<:BlasFloat}}) function frule!!( ::Dual{typeof(LAPACK.getrf!)}, A_dA::Dual{<:AbstractMatrix{P}} -) where {P<:BlasRealFloat} +) where {P<:BlasFloat} _, ipiv, info = LAPACK.getrf!(primal(A_dA)) return _getrf_fwd(A_dA, ipiv, info) end function rrule!!( ::CoDual{typeof(LAPACK.getrf!)}, _A::CoDual{<:AbstractMatrix{P}} -) where {P<:BlasRealFloat} +) where {P<:BlasFloat} A, dA = arrayify(_A) A_copy = copy(A) @@ -28,16 +28,14 @@ end @is_primitive( MinimalCtx, - Tuple{ - typeof(Core.kwcall),NamedTuple,typeof(LAPACK.getrf!),AbstractMatrix{<:BlasRealFloat} - }, + Tuple{typeof(Core.kwcall),NamedTuple,typeof(LAPACK.getrf!),AbstractMatrix{<:BlasFloat}}, ) function frule!!( ::Dual{typeof(Core.kwcall)}, _kwargs::Dual{<:NamedTuple}, ::Dual{typeof(getrf!)}, A_dA::Dual{<:AbstractMatrix{P}}, -) where {P<:BlasRealFloat} +) where {P<:BlasFloat} check = primal(_kwargs).check _, ipiv, info = LAPACK.getrf!(primal(A_dA); check) return _getrf_fwd(A_dA, ipiv, info) @@ -47,7 +45,7 @@ function rrule!!( _kwargs::CoDual{<:NamedTuple}, ::CoDual{typeof(getrf!)}, _A::CoDual{<:AbstractMatrix{P}}, -) where {P<:BlasRealFloat} +) where {P<:BlasFloat} check = _kwargs.x.check A, dA = arrayify(_A) A_copy = copy(A) @@ -505,6 +503,7 @@ end function generate_hand_written_rrule!!_test_cases(rng_ctor, ::Val{:lapack}) rng = rng_ctor(123) Ps = [Float64, Float32] + complexPs = [Float64, Float32, ComplexF64, ComplexF32] bools = [false, true] test_cases = vcat( @@ -516,7 +515,7 @@ function generate_hand_written_rrule!!_test_cases(rng_ctor, ::Val{:lapack}) (false, :stability, nothing, getrf!, A) end end..., - map_prod(bools, Ps) do (check, P) + map_prod(bools, complexPs) do (check, P) As = blas_matrices(rng, P, 5, 5) ipiv = Vector{Int}(undef, 5) return map(As) do A @@ -583,13 +582,33 @@ end function generate_derived_rrule!!_test_cases(rng_ctor, ::Val{:lapack}) rng = rng_ctor(123) + complexPs = [Float64, Float32, ComplexF64, ComplexF32] getrf_wrapper!(x, check) = getrf!(x; check) - test_cases = vcat(map_prod([false, true], [Float64, Float32]) do (check, P) - As = blas_matrices(rng, P, 5, 5) - return map(As) do A - (false, :none, nothing, getrf_wrapper!, A, check) - end - end...) + test_cases = vcat( + # getrf + map_prod([false, true], complexPs) do (check, P) + As = blas_matrices(rng, P, 5, 5) + return map(As) do A + (false, :none, nothing, getrf_wrapper!, A, check) + end + end..., + + # real logdet + map([Float64, Float32]) do P + As = positive_definite_blas_matrices(rng, P, 3) + return map(As) do A + (false, :none, nothing, logdet, A) + end + end..., + + # complex logdet + map(complexPs) do P + As = blas_matrices(rng, P, 3, 3) + return map(As) do A + (false, :none, nothing, real ∘ logdet ∘ complex, A) + end + end..., + ) memory = Any[] return test_cases, memory end diff --git a/src/tangents.jl b/src/tangents.jl index 73a40371b6..5f6c26be51 100644 --- a/src/tangents.jl +++ b/src/tangents.jl @@ -370,6 +370,8 @@ tangent_type(::Type{<:Base.TTY}) = NoTangent tangent_type(::Type{<:IOStream}) = NoTangent +tangent_type(::Type{<:Base.CoreLogging.AbstractLogger}) = NoTangent + function split_union_tuple_type(tangent_types) # Create first split. @@ -1353,6 +1355,7 @@ tangents, but they're unable to check that `increment!!` is correct in an absolu (Union, NoTangent), (UnionAll, NoTangent), (typeof(<:), NoTangent), + (Base.CoreLogging.SimpleLogger, NoTangent), (IOStream(""), NoTangent), ] # Construct test cases containing circular references. These typically require multiple diff --git a/test/fwds_rvs_data.jl b/test/fwds_rvs_data.jl index ebe485ac1c..63f775caf9 100644 --- a/test/fwds_rvs_data.jl +++ b/test/fwds_rvs_data.jl @@ -157,4 +157,9 @@ end @test verify_rdata_value(Ptr{Float64}(), NoRData()) === nothing end end + + @testset "Helpful error messages for misuse of fdata and rdata" begin + @test_throws "Float64 is a type. Perhaps you meant" fdata(Float64) + @test_throws "Float64 is a type. Perhaps you meant" rdata(Float64) + end end diff --git a/test/integration_testing/bijectors/bijectors.jl b/test/integration_testing/bijectors/bijectors.jl index 998c1f34d0..289f97c6d5 100644 --- a/test/integration_testing/bijectors/bijectors.jl +++ b/test/integration_testing/bijectors/bijectors.jl @@ -50,7 +50,7 @@ end 1 0 0 0 0 1 ]), (3, 3)), - # b_binv_test_case(Bijectors.PlanarLayer(3), (3, 3)), + b_binv_test_case(Bijectors.PlanarLayer(3), (3, 3)), b_binv_test_case(Bijectors.RadialLayer(3), 3), b_binv_test_case(Bijectors.Reshape((2, 3), (3, 2)), (2, 3)), b_binv_test_case(Bijectors.Scale(0.2), 3), diff --git a/test/integration_testing/diffeq/Project.toml b/test/integration_testing/diffeq/Project.toml new file mode 100644 index 0000000000..98e015ff48 --- /dev/null +++ b/test/integration_testing/diffeq/Project.toml @@ -0,0 +1,6 @@ +[deps] +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" +OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" +SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1" +StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" \ No newline at end of file diff --git a/test/integration_testing/diffeq/diffeq.jl b/test/integration_testing/diffeq/diffeq.jl new file mode 100644 index 0000000000..b998d5fe41 --- /dev/null +++ b/test/integration_testing/diffeq/diffeq.jl @@ -0,0 +1,63 @@ +using Pkg +Pkg.activate(@__DIR__) +Pkg.develop(; path=joinpath(@__DIR__, "..", "..", "..")) + +using OrdinaryDiffEq, SciMLSensitivity, Mooncake, StableRNGs, Test +using Mooncake.TestUtils: test_rule +using Mooncake: ReverseMode + +# Helper function for Mooncake gradient computation +function mooncake_gradient(f, x) + Mooncake.value_and_gradient!!(Mooncake.build_rrule(f, x), f, x)[2][2] +end + +# Define the ODE function from the original issue +odef(du, u, p, t) = du .= u .* p + +# Define the sensitivity loss function from the original issue +struct senseloss0{T} + sense::T +end + +function (f::senseloss0)(u0p) + prob = ODEProblem{true}(odef, u0p[1:1], (0.0, 1.0), u0p[2:2]) + sum(solve(prob, Tsit5(); abstol=1e-12, reltol=1e-12, saveat=0.1)) +end + +@testset "diffeq" begin + rng = StableRNG(123456) + + # Test parameters from the original issue + u0p = [2.0, 3.0] + + @testset "senseloss0 with InterpolatingAdjoint" begin + sense_func = senseloss0(InterpolatingAdjoint()) + + # First test that the function works + @testset "Function evaluation" begin + result = sense_func(u0p) + @test result isa Real + @test isfinite(result) + end + + # Test Mooncake gradient computation + @testset "mooncake_gradient computation" begin + dup_mc = mooncake_gradient(sense_func, u0p) + @test dup_mc isa Vector + @test length(dup_mc) == 2 + @test all(isfinite, dup_mc) + end + + # Test with Mooncake's test_rule + @testset "test_rule evaluation" begin + test_rule( + rng, + sense_func, + u0p; + is_primitive=false, + unsafe_perturb=true, + mode=ReverseMode, + ) + end + end +end