Use maxthreadid() in TSVI

Author: penelopeysm

src/threadsafe.jl

@@ -13,12 +13,7 @@ function ThreadSafeVarInfo(vi::AbstractVarInfo)
     # fields. This is not good practice --- see
     # https://github.com/TuringLang/DynamicPPL.jl/issues/924 for a full
     # explanation --- but it has worked okay so far.
-    # The use of nthreads()*2 here ensures that threadid() doesn't exceed
-    # the length of the logps array. Ideally, we would use maxthreadid(),
-    # but Mooncake can't differentiate through that. Empirically, nthreads()*2
-    # seems to provide an upper bound to maxthreadid(), so we use that here.
-    # See https://github.com/TuringLang/DynamicPPL.jl/pull/936
-    accs_by_thread = [map(split, getaccs(vi)) for _ in 1:(Threads.nthreads() * 2)]
+    accs_by_thread = [map(split, getaccs(vi)) for _ in 1:Threads.maxthreadid()]
     return ThreadSafeVarInfo(vi, accs_by_thread)
 end
 ThreadSafeVarInfo(vi::ThreadSafeVarInfo) = vi

test/threadsafe.jl

@@ -5,7 +5,7 @@
 
         @test threadsafe_vi.varinfo === vi
         @test threadsafe_vi.accs_by_thread isa Vector{<:DynamicPPL.AccumulatorTuple}
-        @test length(threadsafe_vi.accs_by_thread) == Threads.nthreads() * 2
+        @test length(threadsafe_vi.accs_by_thread) == Threads.maxthreadid()
         expected_accs = DynamicPPL.AccumulatorTuple(
             (DynamicPPL.split(acc) for acc in DynamicPPL.getaccs(vi))...
         )