Adding state batching

Author: murrellb

examples/torus.jl

@@ -39,8 +39,8 @@ sampleX1(n_samples) = Flowfusion.random_literal_cat(n_samples, sigma = T(0.05))[
 n_samples = 500
 
 M = Torus(2)
-P = ManifoldProcess(0.2f0)
-#P = Deterministic()
+P = FProcess(ManifoldProcess(0.2f0), t -> 1-(1-t)^2)
+#P = FProcess(Deterministic(), t -> 1-(1-t)^2)
 
 eta = 0.01
 opt_state = Flux.setup(AdamW(eta = eta, lambda = 0.00001), model)
@@ -57,8 +57,7 @@ for i in 1:iters
     ξ = Guide(Xt, X1)
     #Gradient
     l,g = Flux.withgradient(model) do m
-        #tcloss(P, m(t,tensor(Xt)), ξ, scalefloss(P, t)) #GOING TO HAVE TO ADD GUIDE HERE, AND CHANGE IT TO FLOSS
-        floss(P, m(t,tensor(Xt)), ξ, scalefloss(P, t)) #GOING TO HAVE TO ADD GUIDE HERE, AND CHANGE IT TO FLOSS
+        floss(P, m(t,tensor(Xt)), ξ, scalefloss(P, t))
     end
     #Update
     Flux.update!(opt_state, model, g[1])

src/Flowfusion.jl

@@ -47,7 +47,8 @@ export
     FProcess, 
     floss,
     tcloss,
-    dense
+    dense,
+    batch
 
 
 #Useful for demos etc:

src/bridge.jl

@@ -179,3 +179,20 @@ function apply_tangent_coordinates(Xt::ManifoldState, ξ; retraction_method=defa
     end
     return X̂₁
 end
+
+#As some point we should move these to ForwardBackward.jl. Low prio because they're mostly not needed for other applications.
+tensor_cat(Xs::Vector{T}; dims_from_end = 1) where T = cat(tensor.(Xs)..., dims = ndims(tensor(Xs[1])) - dims_from_end + 1)
+tensor_cat(Xs::Vector{Nothing}) = nothing
+
+"""
+    batch(Xs::Vector{T}; dims_from_end = 1)
+
+Doesn't handle padding. Add option to pad if batching along dims that don't have the same length.
+"""
+batch(Xs::Vector{T}; dims_from_end = 1) where T<:ContinuousState = T(tensor_cat(Xs; dims_from_end))
+batch(Xs::Vector{T}; dims_from_end = 1) where T<:DiscreteState = T(Xs[1].K, tensor_cat(Xs; dims_from_end))
+batch(Xs::Vector{<:ManifoldState{<:M,<:A}}; dims_from_end = 1) where {M, A} = ManifoldState(Xs[1].M, eachslice(tensor_cat(Xs; dims_from_end), dims = Tuple((ndims(Xs[1].state[1])+1:ndims(tensor(Xs[1])))))) #Only tested for rotations.
+batch(Xs::Vector{<:Tuple{Vararg{UState}}}, dims_from_end = 1) = Tuple([batch([x[i] for x in Xs], dims_from_end = dims_from_end) for i in 1:length(Xs[1])])
+
+#Should never move to ForwardBackward.jl
+batch(Xs::Vector{<:MaskedState}; dims_from_end = 1) = MaskedState(batch(unmask.(Xs); dims_from_end), tensor_cat([X.cmask for X in Xs]; dims_from_end), tensor_cat([X.lmask for X in Xs]; dims_from_end))