commit to branch (flow stuff)

Author: AntonOresten

.gitignore

@@ -3,3 +3,4 @@
 *.jl.mem
 /docs/Manifest.toml
 /docs/build/
+Manifest*.toml

Project.toml

@@ -5,6 +5,7 @@ version = "0.1.2"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 ForwardBackward = "e879419d-bb0f-4252-adee-d266c51ac92d"
 Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
@@ -13,15 +14,10 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
 [compat]
 Adapt = "4.1.1"
+FillArrays = "1.13.0"
 ForwardBackward = "0.1.0"
 Manifolds = "0.10.12"
 NNlib = "0.9.27"
 OneHotArrays = "0.2.6"
 StatsBase = "0.34.4"
 julia = "1.9"
-
-[extras]
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-
-[targets]
-test = ["Test"]

src/Flowfusion.jl

@@ -17,9 +17,6 @@ Later:
 =#
 
 
-
-
-
 module Flowfusion
 
 using ForwardBackward, OneHotArrays, Adapt, Manifolds, NNlib
@@ -30,6 +27,8 @@ include("bridge.jl")
 include("loss.jl")
 include("processes.jl")
 
+include("batching.jl")
+
 export
     #Processes not in ForwardBackward.jl
     InterpolatingDiscreteFlow,

src/batching.jl

@@ -1,18 +1,12 @@
-#States
-element(state,seqindex,batchindex) = selectdim(selectdim(state, ndims(state), batchindex), ndims(state)-1, seqindex) 
-element(S::MaskedState, seqindex, batchindex) = element(S.S, seqindex, batchindex)
-element(S::ContinuousState, seqindex, batchindex) = ContinuousState(element(S.state, seqindex, batchindex))
-element(S::ManifoldState, seqindex, batchindex) = ManifoldState(S.M, element(S.state, seqindex, batchindex))
-element(S::DiscreteState, seqindex, batchindex) = DiscreteState(S.K, element(S.state, seqindex, batchindex))
-element(S::Tuple{Vararg{Flowfusion.UState}}, seqindex, batchindex) = element.(S, seqindex, batchindex)
+element(state, seqindex) = selectdim(state, ndims(state), seqindex:seqindex)
+element(state, seqindex, batchindex) = element(selectdim(state, ndims(state), batchindex), seqindex)
 
-#When there isn't a batch dim:
-element(state,seqindex) = selectdim(state, ndims(state), seqindex)
-element(S::MaskedState, seqindex) = element(S.S, seqindex)
-element(S::ContinuousState, seqindex) = ContinuousState(element(S.state, seqindex))
-element(S::ManifoldState, seqindex) = ManifoldState(S.M, element(S.state, seqindex))
-element(S::DiscreteState, seqindex) = DiscreteState(S.K, element(S.state, seqindex))
-element(S::Tuple{Vararg{Flowfusion.UState}}, seqindex) = element.(S, seqindex)
+element(S::MaskedState, inds...) = element(S.S, inds...)
+element(S::ContinuousState, inds...) = ContinuousState(element(S.state, inds...))
+element(S::ManifoldState, inds...) = ManifoldState(S.M, element(S.state, inds...))
+element(S::DiscreteState, inds...) = DiscreteState(S.K, element(S.state, inds...))
+
+element(S::Tuple{Vararg{Flowfusion.UState}}, inds...) = element.(S, inds...)
 
 #Create a "zero" state appropriate for the type. Tricky for manifolds, but we just want rotations working for now I think.
 zerostate(element::T, expandsize...) where T <: ContinuousState = T(similar(tensor(element), size(tensor(element))..., expandsize...) .= 0)

test/Project.toml

@@ -0,0 +1,4 @@
+[deps]
+ForwardBackward = "e879419d-bb0f-4252-adee-d266c51ac92d"
+Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/runtests.jl

@@ -44,7 +44,7 @@ using ForwardBackward
 
     @testset "Bridge, step" begin
 
-        siz = (5,6,7)
+        siz = (5,6)
         XC() = ContinuousState(randn(5, siz...))
         XD() = DiscreteState(5, rand(1:5, siz...))
         MT = Torus(2)