CompatHelper: bump compat for JET in [weakdeps] to 0.10, (keep existing compat)

.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

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"
 

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

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 <https://docs.julialang.org/en/v1/devdocs/ast/#Lowered-form> or <https://docs.julialang.org/en/v1/devdocs/ssair/#Main-SSA-data-structure>. 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.

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()

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),

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

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

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

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),

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"

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