Merge e1bcd9f into d435a0b

TuringLang · Nov 5, 2020 · 979973a · 979973a
2 parents d435a0b + e1bcd9f
commit 979973a
Show file tree

Hide file tree

Showing 4 changed files with 459 additions and 99 deletions.
diff --git a/src/integrator.jl b/src/integrator.jl
@@ -60,13 +60,13 @@ function step(
     h::Hamiltonian,
     z::P,
     n_steps::Int=1;
-    fwd::Bool=n_steps > 0,  # simulate hamiltonian backward when n_steps < 0
+    fwd::Union{Bool,AbstractVector{Bool}}=n_steps > 0,  # simulate hamiltonian backward when n_steps < 0
     full_trajectory::Val{FullTraj} = Val(false)
 ) where {T<:AbstractScalarOrVec{<:AbstractFloat}, P<:PhasePoint, FullTraj}
     n_steps = abs(n_steps)  # to support `n_steps < 0` cases
 
-    ϵ = fwd ? step_size(lf) : -step_size(lf)
-    ϵ = ϵ'
+    ϵ = step_size(lf)
+    ϵ = ifelse.(fwd, ϵ, .-ϵ)'
 
     res = if FullTraj
         Vector{P}(undef, n_steps)