Small tweaks for type stability / inferrability issues

src/prelude.jl

@@ -37,6 +37,9 @@ function structural_iterate(args)
         exclude)
     return iter
 end
+structural_iterate(args::NTuple{N, Union{Real, AbstractFIs}}) where {N} = args
+structural_iterate(args::AbstractArray{T}) where {T <: Union{Real, AbstractFIs}} = args
+structural_iterate(args::T) where {T <: Real} = (args,)
 
 """
     structural_map(f, args)
@@ -46,9 +49,15 @@ often to be used on a function's input/output arguments.
 Currently uses [fmap](https://fluxml.ai/Functors.jl/stable/api/#Functors.fmap) 
 from Functors.jl as a backend.
 """
-function structural_map(f, args...; only_vals = Val{false}())
+function structural_map(f, args...; only_vals = nothing)
+    walk = if only_vals isa Val{true}
+        Functors.StructuralWalk()
+    elseif (only_vals isa Val{false}) || isnothing(only_vals)
+        Functors.DefaultWalk()
+    else
+        error("Unsupported argument only_vals = $only_vals")
+    end
     fmap((args...) -> args[1] isa AbstractArray ? f.(args...) : f(args...), args...;
         cache = nothing,
-        walk = (only_vals isa Val{true} ? Functors.StructuralWalk() :
-                Functors.DefaultWalk()))
+        walk)
 end

src/propagate.jl

@@ -23,13 +23,18 @@ function get_Δs(arg, FIs)
     end
 end
 
-function strip_Δs(arg)
+function strip_Δs(arg; use_dual = Val(true))
     if arg isa StochasticTriple
         # TODO: replace check below with a more robust notion of discreteness.
         if valtype(arg) <: Integer
             return value(arg)
         else
-            return ForwardDiff.Dual{tag(arg)}(value(arg), delta(arg))
+            if use_dual isa Val{true}
+                return ForwardDiff.Dual{tag(arg)}(value(arg), delta(arg))
+            else
+                return StochasticAD.StochasticTriple{tag(arg)}(
+                    value(arg), delta(arg), empty(arg.Δs))
+            end
         end
     else
         return arg
@@ -122,6 +127,7 @@ function propagate(f,
     =#
     out = f(input_args...)
     val = structural_map(value, out)
+    # TODO: what does the only_vals do in the below and why?
     Δs_all = structural_map(Base.Fix2(get_Δs, backendtype(st_rep)), args;
         only_vals = Val{true}())
     # TODO: Coupling approach below needs to handle non-perturbable objects.

src/stochastic_triple.jl

@@ -96,17 +96,21 @@ tag(::ForwardDiff.Dual{T}) where {T} = T
 
 """
     valtype(st::StochasticTriple)
+    valtype(st::Type{<:StochasticTriple})
 
 Get the underlying type of the value tracked by a stochastic triple.
 """
-valtype(::StochasticTriple{T, V}) where {T, V} = V
+valtype(st::StochasticTriple) = valtype(typeof(st))
+valtype(::Type{<:StochasticTriple{T, V}}) where {T, V} = V
 
 """
     backendtype(st::StochasticTriple)
+    backendtype(st::Type{<:StochasticTriple})
 
 Get the backend type of a stochastic triple.
 """
-backendtype(::StochasticTriple{T, V, FIs}) where {T, V, FIs} = FIs
+backendtype(st::StochasticTriple) = backendtype(typeof(st))
+backendtype(::Type{<:StochasticTriple{T, V, FIs}}) where {T, V, FIs} = FIs
 
 """
     smooth_triple(st::StochasticTriple)