Merge branch 'master' into dw/enzyme

TuringLang · Nov 29, 2024 · 2c88d59 · 2c88d59
2 parents 98f2c2b + f4beee9
commit 2c88d59
Show file tree

Hide file tree

Showing 9 changed files with 137 additions and 102 deletions.
diff --git a/.github/workflows/Tests.yml b/.github/workflows/Tests.yml
@@ -9,9 +9,9 @@ on:
 jobs:
   test:
     # Use matrix.test.name here to avoid it taking up the entire window width
-    name: test ${{matrix.test.name}} (${{ matrix.os }}, ${{ matrix.version }}, ${{ matrix.arch }}, ${{ matrix.num_threads }})
-    runs-on: ${{ matrix.os }}
-    continue-on-error: ${{ matrix.version == 'nightly' }}
+    name: test ${{matrix.test.name}} (${{ matrix.runner.os }}, ${{ matrix.runner.version }}, ${{ matrix.runner.arch }}, ${{ matrix.runner.num_threads }})
+    runs-on: ${{ matrix.runner.os }}
+    continue-on-error: ${{ matrix.runner.version == 'pre' }}
 
     strategy:
       fail-fast: false
@@ -35,63 +35,71 @@ jobs:
             args: "mcmc/ess.jl"
           - name: "everything else"
             args: "--skip essential/ad.jl mcmc/gibbs.jl mcmc/hmc.jl mcmc/abstractmcmc.jl mcmc/Inference.jl experimental/gibbs.jl mcmc/ess.jl"
-        version:
-          - '1.10'
-          - '1'
-        os:
-          - ubuntu-latest
-          - windows-latest
-          - macOS-latest
-        arch:
-          - x64
-          - x86
-        num_threads:
-          - 1
-          - 2
-        exclude:
-          # With Windows and macOS, only run x64, 2 threads. We just want to see
-          # some combination work on OSes other than Ubuntu.
-          - os: windows-latest
-            version: '1'
-          - os: macOS-latest
-            version: '1'
-          - os: windows-latest
-            arch: x86
-          - os: macOS-latest
+        runner:
+          # Default
+          - version: '1'
+            os: ubuntu-latest
+            arch: x64
+            num_threads: 1
+          # x86
+          - version: '1'
+            os: ubuntu-latest
             arch: x86
-          - os: windows-latest
             num_threads: 1
-          - os: macOS-latest
+          # Multithreaded
+          - version: '1'
+            os: ubuntu-latest
+            arch: x64
+            num_threads: 2
+          # Windows
+          - version: '1'
+            os: windows-latest
+            arch: x64
+            num_threads: 1
+          # macOS
+          - version: '1'
+            os: macos-latest
+            arch: aarch64
+            num_threads: 1
+          # Minimum supported Julia version
+          - version: 'min'
+            os: ubuntu-latest
+            arch: x64
             num_threads: 1
-          # It's sufficient to test x86 with only Julia 1.10 and 1 thread.
-          - arch: x86
-            version: '1'
-          - arch: x86
+          # Minimum supported Julia version, multithreaded
+          - version: 'min'
+            os: ubuntu-latest
+            arch: x64
             num_threads: 2
+          # Pre-release Julia version
+          - version: 'pre'
+            os: ubuntu-latest
+            arch: x64
+            num_threads: 1
 
     steps:
       - name: Print matrix variables
         run: |
-          echo "OS: ${{ matrix.os }}"
-          echo "Architecture: ${{ matrix.arch }}"
-          echo "Julia version: ${{ matrix.version }}"
-          echo "Number of threads: ${{ matrix.num_threads }}"
+          echo "OS: ${{ matrix.runner.os }}"
+          echo "Architecture: ${{ matrix.runner.arch }}"
+          echo "Julia version: ${{ matrix.runner.version }}"
+          echo "Number of threads: ${{ matrix.runner.num_threads }}"
           echo "Test arguments: ${{ matrix.test.args }}"
       - name: (De)activate coverage analysis
-        run: echo "COVERAGE=${{ matrix.version == '1' && matrix.os == 'ubuntu-latest' && matrix.num_threads == 2 }}" >> "$GITHUB_ENV"
+        run: echo "COVERAGE=${{ matrix.runner.version == '1' && matrix.runner.os == 'ubuntu-latest' && matrix.runner.num_threads == 2 }}" >> "$GITHUB_ENV"
         shell: bash
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2
         with:
-          version: '${{ matrix.version }}'
-          arch: ${{ matrix.arch }}
-      - uses: julia-actions/cache@v1
+          version: '${{ matrix.runner.version }}'
+          arch: ${{ matrix.runner.arch }}
+      - uses: julia-actions/cache@v2
       - uses: julia-actions/julia-buildpkg@v1
       # TODO: Use julia-actions/julia-runtest when test_args are supported
       # Custom calls of Pkg.test tend to miss features such as e.g. adjustments for CompatHelper PRs
       # Ref https://github.com/julia-actions/julia-runtest/pull/73
       - name: Call Pkg.test
-        run: julia --color=yes --inline=yes --depwarn=yes --check-bounds=yes --threads=${{ matrix.num_threads }} --project=@. -e 'import Pkg; Pkg.test(; coverage=parse(Bool, ENV["COVERAGE"]), test_args=ARGS)' -- ${{ matrix.test.args }}
+        run: julia --color=yes --inline=yes --depwarn=yes --check-bounds=yes --threads=${{ matrix.runner.num_threads }} --project=@. -e 'import Pkg; Pkg.test(; coverage=parse(Bool, ENV["COVERAGE"]), test_args=ARGS)' -- ${{ matrix.test.args }}
       - uses: julia-actions/julia-processcoverage@v1
         if: ${{ env.COVERAGE }}
       - uses: codecov/codecov-action@v4

diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "Turing"
 uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
-version = "0.35.0"
+version = "0.35.2"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -56,14 +56,14 @@ AdvancedMH = "0.8"
 AdvancedPS = "0.6.0"
 AdvancedVI = "0.2"
 BangBang = "0.4.2"
-Bijectors = "0.13.6"
+Bijectors = "0.14"
 Compat = "4.15.0"
 DataStructures = "0.18"
 Distributions = "0.23.3, 0.24, 0.25"
 DistributionsAD = "0.6"
 DocStringExtensions = "0.8, 0.9"
 DynamicHMC = "3.4"
-DynamicPPL = "0.29"
+DynamicPPL = "0.29, 0.30.4"
 EllipticalSliceSampling = "0.5, 1, 2"
 ForwardDiff = "0.10.3"
 Libtask = "0.8.8"

diff --git a/src/mcmc/hmc.jl b/src/mcmc/hmc.jl
@@ -181,6 +181,11 @@ function DynamicPPL.initialstep(
             if init_attempt_count == 10
                 @warn "failed to find valid initial parameters in $(init_attempt_count) tries; consider providing explicit initial parameters using the `initial_params` keyword"
             end
+            if init_attempt_count == 1000
+                error(
+                    "failed to find valid initial parameters in $(init_attempt_count) tries. This may indicate an error with the model or AD backend; please open an issue at https://github.com/TuringLang/Turing.jl/issues",
+                )
+            end
 
             # NOTE: This will sample in the unconstrained space.
             vi = last(DynamicPPL.evaluate!!(model, rng, vi, SampleFromUniform()))

diff --git a/test/Project.toml b/test/Project.toml
@@ -46,8 +46,9 @@ Clustering = "0.14, 0.15"
 Distributions = "0.25"
 DistributionsAD = "0.6.3"
 DynamicHMC = "2.1.6, 3.0"
-Enzyme = "0.12, 0.13"
-DynamicPPL = "0.29"
+Enzyme = "0.13"
+DynamicPPL = "0.29, 0.30.2"
+DynamicPPL = "0.29, 0.30.2"
 FiniteDifferences = "0.10.8, 0.11, 0.12"
 ForwardDiff = "0.10.12 - 0.10.32, 0.10"
 HypothesisTests = "0.11"
@@ -72,3 +73,4 @@ StatsBase = "0.33, 0.34"
 StatsFuns = "0.9.5, 1"
 TimerOutputs = "0.5"
 Zygote = "0.5.4, 0.6"
+julia = "1.10"
diff --git a/test/mcmc/abstractmcmc.jl b/test/mcmc/abstractmcmc.jl
@@ -138,23 +138,14 @@ end
                 )
 
                 @testset "inference" begin
-                    if adtype isa AutoReverseDiff &&
-                        model.f === DynamicPPL.TestUtils.demo_assume_index_observe &&
-                        VERSION < v"1.8"
-                        # Ref: https://github.com/TuringLang/DynamicPPL.jl/issues/612
-                        @test_throws UndefRefError sample(
-                            model, sampler_ext, 5_000; sample_kwargs...
-                        )
-                    else
-                        DynamicPPL.TestUtils.test_sampler(
-                            [model],
-                            sampler_ext,
-                            5_000;
-                            rtol=0.2,
-                            sampler_name="AdvancedHMC",
-                            sample_kwargs...,
-                        )
-                    end
+                    DynamicPPL.TestUtils.test_sampler(
+                        [model],
+                        sampler_ext,
+                        5_000;
+                        rtol=0.2,
+                        sampler_name="AdvancedHMC",
+                        sample_kwargs...,
+                    )
                 end
             end
         end

diff --git a/test/mcmc/hmc.jl b/test/mcmc/hmc.jl
@@ -16,7 +16,7 @@ import Random
 using StableRNGs: StableRNG
 using StatsFuns: logistic
 import Mooncake
-using Test: @test, @test_broken, @test_logs, @testset
+using Test: @test, @test_broken, @test_logs, @testset, @test_throws
 using Turing
 
 @testset "Testing hmc.jl with $adbackend" for adbackend in ADUtils.adbackends
@@ -274,6 +274,15 @@ using Turing
         end
     end
 
+    @testset "error for impossible model" begin
+        @model function demo_impossible()
+            x ~ Normal()
+            Turing.@addlogprob! -Inf
+        end
+
+        @test_throws ErrorException sample(demo_impossible(), NUTS(; adtype=adbackend), 5)
+    end
+
     @testset "(partially) issue: #2095" begin
         @model function vector_of_dirichlet(::Type{TV}=Vector{Float64}) where {TV}
             xs = Vector{TV}(undef, 2)

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -1,10 +1,14 @@
 include("test_utils/SelectiveTests.jl")
 using .SelectiveTests: isincluded, parse_args
 using Pkg
+using Random: seed!
 using Test
 using TimerOutputs: TimerOutputs, @timeit
 import Turing
 
+# Fix the global Random.seed for reproducibility.
+seed!(23)
+
 include(pkgdir(Turing) * "/test/test_utils/models.jl")
 include(pkgdir(Turing) * "/test/test_utils/numerical_tests.jl")
 include(pkgdir(Turing) * "/test/test_utils/ad_utils.jl")
@@ -27,6 +31,10 @@ macro timeit_include(path::AbstractString)
 end
 
 @testset "Turing" begin
+    @testset "Test utils" begin
+        @timeit_include("test_utils/test_utils.jl")
+    end
+
     @testset "Aqua" begin
         @timeit_include("Aqua.jl")
     end

diff --git a/test/test_utils/ad_utils.jl b/test/test_utils/ad_utils.jl
@@ -229,44 +229,6 @@ function DynamicPPL.dot_tilde_observe(context::ADTypeCheckContext, sampler, righ
     return logp, vi
 end
 
-# Check that the ADTypeCheckContext works as expected.
-Test.@testset "ADTypeCheckContext" begin
-    Turing.@model test_model() = x ~ Turing.Normal(0, 1)
-    tm = test_model()
-    adtypes = (
-        Turing.AutoForwardDiff(),
-        Turing.AutoReverseDiff(),
-        Turing.AutoZygote(),
-        # TODO: Mooncake
-        # Turing.AutoMooncake(config=nothing),
-    )
-    for actual_adtype in adtypes
-        sampler = Turing.HMC(0.1, 5; adtype=actual_adtype)
-        for expected_adtype in adtypes
-            if (
-                actual_adtype == Turing.AutoForwardDiff() &&
-                expected_adtype == Turing.AutoZygote()
-            )
-                # TODO(mhauru) We are currently unable to check this case.
-                continue
-            end
-            contextualised_tm = DynamicPPL.contextualize(
-                tm, ADTypeCheckContext(expected_adtype, tm.context)
-            )
-            Test.@testset "Expected: $expected_adtype, Actual: $actual_adtype" begin
-                if actual_adtype == expected_adtype
-                    # Check that this does not throw an error.
-                    Turing.sample(contextualised_tm, sampler, 2)
-                else
-                    Test.@test_throws AbstractWrongADBackendError Turing.sample(
-                        contextualised_tm, sampler, 2
-                    )
-                end
-            end
-        end
-    end
-end
-
 # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
 # List of AD backends to test.
 

diff --git a/test/test_utils/test_utils.jl b/test/test_utils/test_utils.jl
@@ -0,0 +1,50 @@
+"""Module for testing the test utils themselves."""
+module TestUtilsTests
+
+using ..ADUtils: ADTypeCheckContext, AbstractWrongADBackendError
+using ForwardDiff: ForwardDiff
+using ReverseDiff: ReverseDiff
+using Test: @test, @testset, @test_throws
+using Turing: Turing
+using Turing: DynamicPPL
+using Zygote: Zygote
+
+# Check that the ADTypeCheckContext works as expected.
+@testset "ADTypeCheckContext" begin
+    Turing.@model test_model() = x ~ Turing.Normal(0, 1)
+    tm = test_model()
+    adtypes = (
+        Turing.AutoForwardDiff(),
+        Turing.AutoReverseDiff(),
+        Turing.AutoZygote(),
+        # TODO: Mooncake
+        # Turing.AutoMooncake(config=nothing),
+    )
+    for actual_adtype in adtypes
+        sampler = Turing.HMC(0.1, 5; adtype=actual_adtype)
+        for expected_adtype in adtypes
+            if (
+                actual_adtype == Turing.AutoForwardDiff() &&
+                expected_adtype == Turing.AutoZygote()
+            )
+                # TODO(mhauru) We are currently unable to check this case.
+                continue
+            end
+            contextualised_tm = DynamicPPL.contextualize(
+                tm, ADTypeCheckContext(expected_adtype, tm.context)
+            )
+            @testset "Expected: $expected_adtype, Actual: $actual_adtype" begin
+                if actual_adtype == expected_adtype
+                    # Check that this does not throw an error.
+                    Turing.sample(contextualised_tm, sampler, 2)
+                else
+                    @test_throws AbstractWrongADBackendError Turing.sample(
+                        contextualised_tm, sampler, 2
+                    )
+                end
+            end
+        end
+    end
+end
+
+end