fix JET tests (#479)

Author: nsiccha

src/integrator.jl

@@ -89,14 +89,14 @@ Leapfrog integrator with randomly "jittered" step size `ϵ` for every trajectory
 $(TYPEDFIELDS)
 
 # Description
-This is the same as `LeapFrog`(@ref) but with a "jittered" step size. This means 
-that at the beginning of each trajectory we sample a step size `ϵ` by adding or 
-subtracting from the nominal/base step size `ϵ0` some random proportion of `ϵ0`, 
+This is the same as `LeapFrog`(@ref) but with a "jittered" step size. This means
+that at the beginning of each trajectory we sample a step size `ϵ` by adding or
+subtracting from the nominal/base step size `ϵ0` some random proportion of `ϵ0`,
 with the proportion specified by `jitter`, i.e. `ϵ = ϵ0 - jitter * ϵ0 * rand()`.
 p
 Jittering might help alleviate issues related to poor interactions with a fixed step size:
-- In regions with high "curvature" the current choice of step size might mean over-shoot 
-  leading to almost all steps being rejected. Randomly sampling the step size at the 
+- In regions with high "curvature" the current choice of step size might mean over-shoot
+  leading to almost all steps being rejected. Randomly sampling the step size at the
   beginning of the trajectories can therefore increase the probability of escaping such
   high-curvature regions.
 - Exact periodicity of the simulated trajectories might occur, i.e. you might be so
@@ -168,7 +168,7 @@ $(TYPEDFIELDS)
 
 # Description
 
-Tempering can potentially allow greater exploration of the posterior, e.g. 
+Tempering can potentially allow greater exploration of the posterior, e.g.
 in a multi-modal posterior jumps between the modes can be more likely to occur.
 """
 struct TemperedLeapfrog{FT<:AbstractFloat,T<:AbstractScalarOrVec{FT}} <: AbstractLeapfrog{T}
@@ -226,10 +226,7 @@ function step(
     ϵ = fwd ? step_size(lf) : -step_size(lf)
     ϵ = ϵ'
 
-    res = nothing
-    if FullTraj
-        res = Vector{P}(undef, n_steps)
-    end
+    res = FullTraj ? Vector{P}(undef, n_steps) : nothing
 
     (; θ, r) = z
     (; value, gradient) = z.ℓπ
@@ -249,12 +246,12 @@ function step(
         # Create a new phase point by caching the logdensity and gradient
         z = phasepoint(h, θ, r; ℓπ=DualValue(value, gradient))
         # Update result
-        if FullTraj
+        if !isnothing(res)
             res[i] = z
         end
         if !isfinite(z)
             # Remove undef
-            if FullTraj
+            if !isnothing(res)
                 resize!(res, i)
             end
             break

src/riemannian/integrator.jl

@@ -11,7 +11,7 @@ Generalized leapfrog integrator with fixed step size `ϵ`.
 $(TYPEDFIELDS)
 
 
-## References 
+## References
 
 1. Girolami, Mark, and Ben Calderhead. "Riemann manifold Langevin and Hamiltonian Monte Carlo methods." Journal of the Royal Statistical Society Series B: Statistical Methodology 73, no. 2 (2011): 123-214.
 """
@@ -77,7 +77,7 @@ end
 
 # TODO(Kai) make sure vectorization works
 # TODO(Kai) check if tempering is valid
-# TODO(Kai) abstract out the 3 main steps and merge with `step` in `integrator.jl` 
+# TODO(Kai) abstract out the 3 main steps and merge with `step` in `integrator.jl`
 function step(
     lf::GeneralizedLeapfrog{T},
     h::Hamiltonian{<:DenseRiemannianMetric},