Package evaluation of AutoEncoderToolkit on Julia 1.12.0-DEV.1805 (a080deafdd*) started at 2025-03-25T00:31:48.775


################################################################################
# Set-up
#

Installing PkgEval dependencies (TestEnv)...

Set-up completed after 9.1s


################################################################################
# Installation
#

Installing AutoEncoderToolkit...
   Resolving package versions...
    Updating `~/.julia/environments/v1.12/Project.toml`
  [1575904b] + AutoEncoderToolkit v0.1.2
    Updating `~/.julia/environments/v1.12/Manifest.toml`
  [47edcb42] + ADTypes v1.14.0
  [621f4979] + AbstractFFTs v1.5.0
  [1520ce14] + AbstractTrees v0.4.5
  [7d9f7c33] + Accessors v0.1.42
  [79e6a3ab] + Adapt v4.3.0
  [66dad0bd] + AliasTables v1.1.3
  [dce04be8] + ArgCheck v2.5.0
  [4fba245c] + ArrayInterface v7.18.0
  [a9b6321e] + Atomix v1.1.1
  [1575904b] + AutoEncoderToolkit v0.1.2
  [198e06fe] + BangBang v0.4.4
  [9718e550] + Baselet v0.1.1
  [e2ed5e7c] + Bijections v0.1.9
  [fa961155] + CEnum v0.5.0
  [082447d4] + ChainRules v1.72.3
  [d360d2e6] + ChainRulesCore v1.25.1
  [aaaa29a8] + Clustering v0.15.8
  [861a8166] + Combinatorics v1.0.2
  [38540f10] + CommonSolve v0.2.4
  [bbf7d656] + CommonSubexpressions v0.3.1
  [f70d9fcc] + CommonWorldInvalidations v1.0.0
  [34da2185] + Compat v4.16.0
  [b152e2b5] + CompositeTypes v0.1.4
  [a33af91c] + CompositionsBase v0.1.2
  [2569d6c7] + ConcreteStructs v0.2.3
  [187b0558] + ConstructionBase v1.5.8
  [6add18c4] + ContextVariablesX v0.1.3
  [9a962f9c] + DataAPI v1.16.0
  [864edb3b] + DataStructures v0.18.22
  [e2d170a0] + DataValueInterfaces v1.0.0
  [244e2a9f] + DefineSingletons v0.1.2
  [8bb1440f] + DelimitedFiles v1.9.1
  [163ba53b] + DiffResults v1.1.0
  [b552c78f] + DiffRules v1.15.1
  [b4f34e82] + Distances v0.10.12
  [31c24e10] + Distributions v0.25.118
  [ffbed154] + DocStringExtensions v0.9.3
  [5b8099bc] + DomainSets v0.7.15
  [7c1d4256] + DynamicPolynomials v0.6.1
  [da5c29d0] + EllipsisNotation v1.8.0
  [4e289a0a] + EnumX v1.0.4
  [e2ba6199] + ExprTools v0.1.10
  [55351af7] + ExproniconLite v0.10.14
  [cc61a311] + FLoops v0.2.2
  [b9860ae5] + FLoopsBase v0.1.1
  [1a297f60] + FillArrays v1.13.0
⌅ [587475ba] + Flux v0.14.25
  [1fa38f19] + Format v1.3.7
  [f6369f11] + ForwardDiff v0.10.38
  [069b7b12] + FunctionWrappers v1.1.3
  [77dc65aa] + FunctionWrappersWrappers v0.1.3
⌅ [d9f16b24] + Functors v0.4.12
  [0c68f7d7] + GPUArrays v11.2.2
  [46192b85] + GPUArraysCore v0.2.0
  [076d061b] + HashArrayMappedTries v0.2.0
  [34004b35] + HypergeometricFunctions v0.3.28
  [7869d1d1] + IRTools v0.4.14
  [615f187c] + IfElse v0.1.1
  [22cec73e] + InitialValues v0.3.1
  [18e54dd8] + IntegerMathUtils v0.1.2
  [8197267c] + IntervalSets v0.7.10
  [3587e190] + InverseFunctions v0.1.17
  [92d709cd] + IrrationalConstants v0.2.4
  [82899510] + IteratorInterfaceExtensions v1.0.0
  [692b3bcd] + JLLWrappers v1.7.0
  [ae98c720] + Jieko v0.2.1
  [b14d175d] + JuliaVariables v0.2.4
  [63c18a36] + KernelAbstractions v0.9.34
  [929cbde3] + LLVM v9.2.0
  [b964fa9f] + LaTeXStrings v1.4.0
  [23fbe1c1] + Latexify v0.16.6
  [2ab3a3ac] + LogExpFunctions v0.3.29
  [c2834f40] + MLCore v1.0.0
⌃ [7e8f7934] + MLDataDevices v1.5.3
  [d8e11817] + MLStyle v0.4.17
  [f1d291b0] + MLUtils v0.4.8
  [1914dd2f] + MacroTools v0.5.15
  [128add7d] + MicroCollections v0.2.0
  [e1d29d7a] + Missings v1.2.0
  [2e0e35c7] + Moshi v0.3.5
  [102ac46a] + MultivariatePolynomials v0.5.7
  [d8a4904e] + MutableArithmetics v1.6.4
  [872c559c] + NNlib v0.9.29
  [77ba4419] + NaNMath v1.1.2
  [71a1bf82] + NameResolution v0.1.5
  [b8a86587] + NearestNeighbors v0.4.21
  [6fe1bfb0] + OffsetArrays v1.16.0
  [0b1bfda6] + OneHotArrays v0.2.6
⌅ [3bd65402] + Optimisers v0.3.4
  [bac558e1] + OrderedCollections v1.8.0
  [90014a1f] + PDMats v0.11.32
  [aea7be01] + PrecompileTools v1.2.1
  [21216c6a] + Preferences v1.4.3
  [8162dcfd] + PrettyPrint v0.2.0
  [27ebfcd6] + Primes v0.5.7
  [33c8b6b6] + ProgressLogging v0.1.4
  [43287f4e] + PtrArrays v1.3.0
  [1fd47b50] + QuadGK v2.11.2
  [c1ae055f] + RealDot v0.1.0
  [3cdcf5f2] + RecipesBase v1.3.4
  [731186ca] + RecursiveArrayTools v3.31.1
  [189a3867] + Reexport v1.2.2
  [ae029012] + Requires v1.3.1
  [79098fc4] + Rmath v0.8.0
  [7e49a35a] + RuntimeGeneratedFunctions v0.5.13
  [0bca4576] + SciMLBase v2.79.0
  [c0aeaf25] + SciMLOperators v0.3.13
  [53ae85a6] + SciMLStructures v1.7.0
  [7e506255] + ScopedValues v1.3.0
  [efcf1570] + Setfield v1.1.2
  [605ecd9f] + ShowCases v0.1.0
  [699a6c99] + SimpleTraits v0.9.4
  [a2af1166] + SortingAlgorithms v1.2.1
  [dc90abb0] + SparseInverseSubset v0.1.2
  [276daf66] + SpecialFunctions v2.5.0
  [171d559e] + SplittablesBase v0.1.15
  [aedffcd0] + Static v1.2.0
  [0d7ed370] + StaticArrayInterface v1.8.0
  [90137ffa] + StaticArrays v1.9.13
  [1e83bf80] + StaticArraysCore v1.4.3
  [10745b16] + Statistics v1.11.1
  [82ae8749] + StatsAPI v1.7.0
  [2913bbd2] + StatsBase v0.34.4
  [4c63d2b9] + StatsFuns v1.3.2
  [09ab397b] + StructArrays v0.7.0
  [2efcf032] + SymbolicIndexingInterface v0.3.38
  [19f23fe9] + SymbolicLimits v0.2.2
  [d1185830] + SymbolicUtils v3.22.0
  [0c5d862f] + Symbolics v6.33.1
  [3783bdb8] + TableTraits v1.0.1
  [bd369af6] + Tables v1.12.0
  [ed4db957] + TaskLocalValues v0.1.2
⌅ [b36ab563] + TaylorDiff v0.2.5
  [8ea1fca8] + TermInterface v2.0.0
  [a759f4b9] + TimerOutputs v0.5.28
  [28d57a85] + Transducers v0.4.84
  [a7c27f48] + Unityper v0.1.6
  [013be700] + UnsafeAtomics v0.3.0
⌅ [e88e6eb3] + Zygote v0.6.75
  [700de1a5] + ZygoteRules v0.2.7
  [dad2f222] + LLVMExtra_jll v0.0.35+0
  [efe28fd5] + OpenSpecFun_jll v0.5.6+0
  [f50d1b31] + Rmath_jll v0.5.1+0
  [0dad84c5] + ArgTools v1.1.2
  [56f22d72] + Artifacts v1.11.0
  [2a0f44e3] + Base64 v1.11.0
  [ade2ca70] + Dates v1.11.0
  [8ba89e20] + Distributed v1.11.0
  [f43a241f] + Downloads v1.6.0
  [7b1f6079] + FileWatching v1.11.0
  [9fa8497b] + Future v1.11.0
  [b77e0a4c] + InteractiveUtils v1.11.0
  [dc6e5ff7] + JuliaSyntaxHighlighting v1.12.0
  [4af54fe1] + LazyArtifacts v1.11.0
  [b27032c2] + LibCURL v0.6.4
  [76f85450] + LibGit2 v1.11.0
  [8f399da3] + Libdl v1.11.0
  [37e2e46d] + LinearAlgebra v1.11.0
  [56ddb016] + Logging v1.11.0
  [d6f4376e] + Markdown v1.11.0
  [a63ad114] + Mmap v1.11.0
  [ca575930] + NetworkOptions v1.2.0
  [44cfe95a] + Pkg v1.12.0
  [de0858da] + Printf v1.11.0
  [9a3f8284] + Random v1.11.0
  [ea8e919c] + SHA v0.7.0
  [9e88b42a] + Serialization v1.11.0
  [6462fe0b] + Sockets v1.11.0
  [2f01184e] + SparseArrays v1.12.0
  [f489334b] + StyledStrings v1.11.0
  [4607b0f0] + SuiteSparse
  [fa267f1f] + TOML v1.0.3
  [a4e569a6] + Tar v1.10.0
  [8dfed614] + Test v1.11.0
  [cf7118a7] + UUIDs v1.11.0
  [4ec0a83e] + Unicode v1.11.0
  [e66e0078] + CompilerSupportLibraries_jll v1.2.0+0
  [deac9b47] + LibCURL_jll v8.6.0+0
  [e37daf67] + LibGit2_jll v1.8.0+0
  [29816b5a] + LibSSH2_jll v1.11.0+1
  [c8ffd9c3] + MbedTLS_jll v2.28.6+1
  [14a3606d] + MozillaCACerts_jll v2024.11.26
  [4536629a] + OpenBLAS_jll v0.3.28+3
  [05823500] + OpenLibm_jll v0.8.1+3
  [bea87d4a] + SuiteSparse_jll v7.8.0+1
  [83775a58] + Zlib_jll v1.3.1+1
  [8e850b90] + libblastrampoline_jll v5.11.2+0
  [8e850ede] + nghttp2_jll v1.63.0+1
  [3f19e933] + p7zip_jll v17.5.0+1
        Info Packages marked with ⌃ and ⌅ have new versions available. Those with ⌃ may be upgradable, but those with ⌅ are restricted by compatibility constraints from upgrading. To see why use `status --outdated -m`

Installation completed after 6.16s


################################################################################
# Precompilation
#

Precompiling PkgEval dependencies...

Precompiling package dependencies...

Precompilation completed after 642.1s


################################################################################
# Testing
#

     Testing AutoEncoderToolkit
      Status `/tmp/jl_TycZd8/Project.toml`
  [1575904b] AutoEncoderToolkit v0.1.2
  [d360d2e6] ChainRulesCore v1.25.1
  [aaaa29a8] Clustering v0.15.8
  [2569d6c7] ConcreteStructs v0.2.3
  [b4f34e82] Distances v0.10.12
  [31c24e10] Distributions v0.25.118
  [da5c29d0] EllipsisNotation v1.8.0
⌅ [587475ba] Flux v0.14.25
  [f6369f11] ForwardDiff v0.10.38
  [b8a86587] NearestNeighbors v0.4.21
  [2913bbd2] StatsBase v0.34.4
⌅ [b36ab563] TaylorDiff v0.2.5
⌅ [e88e6eb3] Zygote v0.6.75
  [37e2e46d] LinearAlgebra v1.11.0
  [44cfe95a] Pkg v1.12.0
  [9a3f8284] Random v1.11.0
  [8dfed614] Test v1.11.0
      Status `/tmp/jl_TycZd8/Manifest.toml`
  [47edcb42] ADTypes v1.14.0
  [621f4979] AbstractFFTs v1.5.0
  [1520ce14] AbstractTrees v0.4.5
  [7d9f7c33] Accessors v0.1.42
  [79e6a3ab] Adapt v4.3.0
  [66dad0bd] AliasTables v1.1.3
  [dce04be8] ArgCheck v2.5.0
  [4fba245c] ArrayInterface v7.18.0
  [a9b6321e] Atomix v1.1.1
  [1575904b] AutoEncoderToolkit v0.1.2
  [198e06fe] BangBang v0.4.4
  [9718e550] Baselet v0.1.1
  [e2ed5e7c] Bijections v0.1.9
  [fa961155] CEnum v0.5.0
  [082447d4] ChainRules v1.72.3
  [d360d2e6] ChainRulesCore v1.25.1
  [aaaa29a8] Clustering v0.15.8
  [861a8166] Combinatorics v1.0.2
  [38540f10] CommonSolve v0.2.4
  [bbf7d656] CommonSubexpressions v0.3.1
  [f70d9fcc] CommonWorldInvalidations v1.0.0
  [34da2185] Compat v4.16.0
  [b152e2b5] CompositeTypes v0.1.4
  [a33af91c] CompositionsBase v0.1.2
  [2569d6c7] ConcreteStructs v0.2.3
  [187b0558] ConstructionBase v1.5.8
  [6add18c4] ContextVariablesX v0.1.3
  [9a962f9c] DataAPI v1.16.0
  [864edb3b] DataStructures v0.18.22
  [e2d170a0] DataValueInterfaces v1.0.0
  [244e2a9f] DefineSingletons v0.1.2
  [8bb1440f] DelimitedFiles v1.9.1
  [163ba53b] DiffResults v1.1.0
  [b552c78f] DiffRules v1.15.1
  [b4f34e82] Distances v0.10.12
  [31c24e10] Distributions v0.25.118
  [ffbed154] DocStringExtensions v0.9.3
  [5b8099bc] DomainSets v0.7.15
  [7c1d4256] DynamicPolynomials v0.6.1
  [da5c29d0] EllipsisNotation v1.8.0
  [4e289a0a] EnumX v1.0.4
  [e2ba6199] ExprTools v0.1.10
  [55351af7] ExproniconLite v0.10.14
  [cc61a311] FLoops v0.2.2
  [b9860ae5] FLoopsBase v0.1.1
  [1a297f60] FillArrays v1.13.0
⌅ [587475ba] Flux v0.14.25
  [1fa38f19] Format v1.3.7
  [f6369f11] ForwardDiff v0.10.38
  [069b7b12] FunctionWrappers v1.1.3
  [77dc65aa] FunctionWrappersWrappers v0.1.3
⌅ [d9f16b24] Functors v0.4.12
  [0c68f7d7] GPUArrays v11.2.2
  [46192b85] GPUArraysCore v0.2.0
  [076d061b] HashArrayMappedTries v0.2.0
  [34004b35] HypergeometricFunctions v0.3.28
  [7869d1d1] IRTools v0.4.14
  [615f187c] IfElse v0.1.1
  [22cec73e] InitialValues v0.3.1
  [18e54dd8] IntegerMathUtils v0.1.2
  [8197267c] IntervalSets v0.7.10
  [3587e190] InverseFunctions v0.1.17
  [92d709cd] IrrationalConstants v0.2.4
  [82899510] IteratorInterfaceExtensions v1.0.0
  [692b3bcd] JLLWrappers v1.7.0
  [ae98c720] Jieko v0.2.1
  [b14d175d] JuliaVariables v0.2.4
  [63c18a36] KernelAbstractions v0.9.34
  [929cbde3] LLVM v9.2.0
  [b964fa9f] LaTeXStrings v1.4.0
  [23fbe1c1] Latexify v0.16.6
  [2ab3a3ac] LogExpFunctions v0.3.29
  [c2834f40] MLCore v1.0.0
⌃ [7e8f7934] MLDataDevices v1.5.3
  [d8e11817] MLStyle v0.4.17
  [f1d291b0] MLUtils v0.4.8
  [1914dd2f] MacroTools v0.5.15
  [128add7d] MicroCollections v0.2.0
  [e1d29d7a] Missings v1.2.0
  [2e0e35c7] Moshi v0.3.5
  [102ac46a] MultivariatePolynomials v0.5.7
  [d8a4904e] MutableArithmetics v1.6.4
  [872c559c] NNlib v0.9.29
  [77ba4419] NaNMath v1.1.2
  [71a1bf82] NameResolution v0.1.5
  [b8a86587] NearestNeighbors v0.4.21
  [6fe1bfb0] OffsetArrays v1.16.0
  [0b1bfda6] OneHotArrays v0.2.6
⌅ [3bd65402] Optimisers v0.3.4
  [bac558e1] OrderedCollections v1.8.0
  [90014a1f] PDMats v0.11.32
  [aea7be01] PrecompileTools v1.2.1
  [21216c6a] Preferences v1.4.3
  [8162dcfd] PrettyPrint v0.2.0
  [27ebfcd6] Primes v0.5.7
  [33c8b6b6] ProgressLogging v0.1.4
  [43287f4e] PtrArrays v1.3.0
  [1fd47b50] QuadGK v2.11.2
  [c1ae055f] RealDot v0.1.0
  [3cdcf5f2] RecipesBase v1.3.4
  [731186ca] RecursiveArrayTools v3.31.1
  [189a3867] Reexport v1.2.2
  [ae029012] Requires v1.3.1
  [79098fc4] Rmath v0.8.0
  [7e49a35a] RuntimeGeneratedFunctions v0.5.13
  [0bca4576] SciMLBase v2.79.0
  [c0aeaf25] SciMLOperators v0.3.13
  [53ae85a6] SciMLStructures v1.7.0
  [7e506255] ScopedValues v1.3.0
  [efcf1570] Setfield v1.1.2
  [605ecd9f] ShowCases v0.1.0
  [699a6c99] SimpleTraits v0.9.4
  [a2af1166] SortingAlgorithms v1.2.1
  [dc90abb0] SparseInverseSubset v0.1.2
  [276daf66] SpecialFunctions v2.5.0
  [171d559e] SplittablesBase v0.1.15
  [aedffcd0] Static v1.2.0
  [0d7ed370] StaticArrayInterface v1.8.0
  [90137ffa] StaticArrays v1.9.13
  [1e83bf80] StaticArraysCore v1.4.3
  [10745b16] Statistics v1.11.1
  [82ae8749] StatsAPI v1.7.0
  [2913bbd2] StatsBase v0.34.4
  [4c63d2b9] StatsFuns v1.3.2
  [09ab397b] StructArrays v0.7.0
  [2efcf032] SymbolicIndexingInterface v0.3.38
  [19f23fe9] SymbolicLimits v0.2.2
  [d1185830] SymbolicUtils v3.22.0
  [0c5d862f] Symbolics v6.33.1
  [3783bdb8] TableTraits v1.0.1
  [bd369af6] Tables v1.12.0
  [ed4db957] TaskLocalValues v0.1.2
⌅ [b36ab563] TaylorDiff v0.2.5
  [8ea1fca8] TermInterface v2.0.0
  [a759f4b9] TimerOutputs v0.5.28
  [28d57a85] Transducers v0.4.84
  [a7c27f48] Unityper v0.1.6
  [013be700] UnsafeAtomics v0.3.0
⌅ [e88e6eb3] Zygote v0.6.75
  [700de1a5] ZygoteRules v0.2.7
  [dad2f222] LLVMExtra_jll v0.0.35+0
  [efe28fd5] OpenSpecFun_jll v0.5.6+0
  [f50d1b31] Rmath_jll v0.5.1+0
  [0dad84c5] ArgTools v1.1.2
  [56f22d72] Artifacts v1.11.0
  [2a0f44e3] Base64 v1.11.0
  [ade2ca70] Dates v1.11.0
  [8ba89e20] Distributed v1.11.0
  [f43a241f] Downloads v1.6.0
  [7b1f6079] FileWatching v1.11.0
  [9fa8497b] Future v1.11.0
  [b77e0a4c] InteractiveUtils v1.11.0
  [dc6e5ff7] JuliaSyntaxHighlighting v1.12.0
  [4af54fe1] LazyArtifacts v1.11.0
  [b27032c2] LibCURL v0.6.4
  [76f85450] LibGit2 v1.11.0
  [8f399da3] Libdl v1.11.0
  [37e2e46d] LinearAlgebra v1.11.0
  [56ddb016] Logging v1.11.0
  [d6f4376e] Markdown v1.11.0
  [a63ad114] Mmap v1.11.0
  [ca575930] NetworkOptions v1.2.0
  [44cfe95a] Pkg v1.12.0
  [de0858da] Printf v1.11.0
  [9a3f8284] Random v1.11.0
  [ea8e919c] SHA v0.7.0
  [9e88b42a] Serialization v1.11.0
  [6462fe0b] Sockets v1.11.0
  [2f01184e] SparseArrays v1.12.0
  [f489334b] StyledStrings v1.11.0
  [4607b0f0] SuiteSparse
  [fa267f1f] TOML v1.0.3
  [a4e569a6] Tar v1.10.0
  [8dfed614] Test v1.11.0
  [cf7118a7] UUIDs v1.11.0
  [4ec0a83e] Unicode v1.11.0
  [e66e0078] CompilerSupportLibraries_jll v1.2.0+0
  [deac9b47] LibCURL_jll v8.6.0+0
  [e37daf67] LibGit2_jll v1.8.0+0
  [29816b5a] LibSSH2_jll v1.11.0+1
  [c8ffd9c3] MbedTLS_jll v2.28.6+1
  [14a3606d] MozillaCACerts_jll v2024.11.26
  [4536629a] OpenBLAS_jll v0.3.28+3
  [05823500] OpenLibm_jll v0.8.1+3
  [bea87d4a] SuiteSparse_jll v7.8.0+1
  [83775a58] Zlib_jll v1.3.1+1
  [8e850b90] libblastrampoline_jll v5.11.2+0
  [8e850ede] nghttp2_jll v1.63.0+1
  [3f19e933] p7zip_jll v17.5.0+1
        Info Packages marked with ⌃ and ⌅ have new versions available. Those with ⌃ may be upgradable, but those with ⌅ are restricted by compatibility constraints from upgrading.
     Testing Running tests...
Testing AutoEncoderToolkit.jl

CUDA is not available - skipping CUDA tests.


Testing utils module:


All tests passed!


Testin Encoders...


All tests passed!


Testing decoders...

We changed the last layer number of µ_neurons to match the input dimension
We changed the last layer number of logσ_neurons to match the input dimension
We changed the last layer number of µ_neurons to match the input dimension
We changed the last layer number of σ_neurons to match the input dimension
We changed the last layer number of µ_neurons to match the input dimension
We changed the last layer number of logσ_neurons to match the input dimension
We changed the last layer number of µ_neurons to match the input dimension
We changed the last layer number of σ_neurons to match the input dimension

All tests passed!


Testing AEs module...


All tests passed!


Testing VAEs module...

Generating synthetic data...
Defining layers structure...
Defining encoders...
Defining decoders...
with same input as output: Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:334
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [3] (::Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [4] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##reparameterize#1", Bool, Type{Float32}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [5] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [6] reparameterize
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:157 [inlined]
    [7] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.reparameterize), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:log,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{log::Bool}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##_#5", Bool, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [9] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [10] #loss#6
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:344 [inlined]
   [11] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, Matrix{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [12] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:313 [inlined]
   [13] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, Matrix{Float32}}, Any}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [14] (::Zygote.Pullback{Tuple{var"#45#46"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#45#46"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, Matrix{Float32}}, Any}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [15] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#45#46"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#45#46"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, Matrix{Float32}}, Any}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [16] gradient(f::Function, args::AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [17] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:324
   [18] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [19] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:335 [inlined]
   [20] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [21] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:335 [inlined]
   [22] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [23] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [24] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [25] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:30 [inlined]
   [26] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [27] top-level scope
      @ none:6
   [28] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [29] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
   [30] _start()
      @ Base ./client.jl:558
with different input and output: Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:339
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [3] (::Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [4] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##reparameterize#1", Bool, Type{Float32}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [5] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [6] reparameterize
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:157 [inlined]
    [7] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.reparameterize), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:log,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{log::Bool}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##_#5", Bool, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [9] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [10] #loss#7
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:438 [inlined]
   [11] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [12] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:406 [inlined]
   [13] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [14] (::Zygote.Pullback{Tuple{var"#47#48"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [15] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#47#48"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [16] gradient(f::Function, args::AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [17] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:324
   [18] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [19] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:340 [inlined]
   [20] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [21] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:340 [inlined]
   [22] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [23] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [24] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [25] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:30 [inlined]
   [26] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [27] top-level scope
      @ none:6
   [28] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [29] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
   [30] _start()
      @ Base ./client.jl:558
with same input as output: Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:334
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [3] (::Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [4] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##reparameterize#1", Bool, Type{Float32}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [5] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [6] reparameterize
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:157 [inlined]
    [7] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.reparameterize), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:log,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{log::Bool}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##_#5", Bool, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [9] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [10] #loss#6
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:344 [inlined]
   [11] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}, Matrix{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [12] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:313 [inlined]
   [13] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}, Matrix{Float32}}, Any}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [14] (::Zygote.Pullback{Tuple{var"#45#46"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#45#46"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}, Matrix{Float32}}, Any}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [15] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#45#46"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#45#46"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}, Matrix{Float32}}, Any}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [16] gradient(f::Function, args::AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [17] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:324
   [18] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [19] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:335 [inlined]
   [20] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [21] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:335 [inlined]
   [22] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [23] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [24] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [25] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:30 [inlined]
   [26] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [27] top-level scope
      @ none:6
   [28] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [29] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
   [30] _start()
      @ Base ./client.jl:558
with different input and output: Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:339
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [3] (::Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [4] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##reparameterize#1", Bool, Type{Float32}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [5] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [6] reparameterize
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:157 [inlined]
    [7] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.reparameterize), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:log,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{log::Bool}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##_#5", Bool, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [9] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [10] #loss#7
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:438 [inlined]
   [11] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [12] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:406 [inlined]
   [13] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [14] (::Zygote.Pullback{Tuple{var"#47#48"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [15] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#47#48"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [16] gradient(f::Function, args::AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [17] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:324
   [18] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [19] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:340 [inlined]
   [20] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [21] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:340 [inlined]
   [22] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [23] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [24] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [25] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:30 [inlined]
   [26] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [27] top-level scope
      @ none:6
   [28] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [29] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
   [30] _start()
      @ Base ./client.jl:558
with same input as output: Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:334
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [3] (::Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [4] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##reparameterize#1", Bool, Type{Float32}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [5] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [6] reparameterize
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:157 [inlined]
    [7] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.reparameterize), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:log,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{log::Bool}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##_#5", Bool, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{μ::Matrix{Float32}, σ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [9] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{μ::Matrix{Float32}, σ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [10] #loss#6
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:344 [inlined]
   [11] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [12] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:313 [inlined]
   [13] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}, Matrix{Float32}}, Any}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [14] (::Zygote.Pullback{Tuple{var"#45#46"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#45#46"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}, Matrix{Float32}}, Any}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [15] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#45#46"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#45#46"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}, Matrix{Float32}}, Any}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [16] gradient(f::Function, args::AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [17] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:324
   [18] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [19] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:335 [inlined]
   [20] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [21] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:335 [inlined]
   [22] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [23] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [24] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [25] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:30 [inlined]
   [26] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [27] top-level scope
      @ none:6
   [28] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [29] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
   [30] _start()
      @ Base ./client.jl:558
with different input and output: Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:339
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [3] (::Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [4] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##reparameterize#1", Bool, Type{Float32}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [5] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [6] reparameterize
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:157 [inlined]
    [7] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.reparameterize), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:log,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{log::Bool}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##_#5", Bool, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{μ::Matrix{Float32}, σ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [9] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{μ::Matrix{Float32}, σ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [10] #loss#7
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:438 [inlined]
   [11] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [12] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:406 [inlined]
   [13] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [14] (::Zygote.Pullback{Tuple{var"#47#48"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [15] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#47#48"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [16] gradient(f::Function, args::AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [17] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:324
   [18] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [19] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:340 [inlined]
   [20] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [21] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:340 [inlined]
   [22] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [23] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [24] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [25] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:30 [inlined]
   [26] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [27] top-level scope
      @ none:6
   [28] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [29] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
   [30] _start()
      @ Base ./client.jl:558
with same input as output: Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:334
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [3] (::Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [4] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##reparameterize#1", Bool, Type{Float32}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [5] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [6] reparameterize
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:157 [inlined]
    [7] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.reparameterize), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:log,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{log::Bool}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##_#5", Bool, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{μ::Matrix{Float32}, σ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [9] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{μ::Matrix{Float32}, σ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [10] #loss#6
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:344 [inlined]
   [11] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [12] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:313 [inlined]
   [13] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}, Matrix{Float32}}, Any}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [14] (::Zygote.Pullback{Tuple{var"#45#46"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#45#46"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}, Matrix{Float32}}, Any}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [15] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#45#46"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#45#46"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}, Matrix{Float32}}, Any}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [16] gradient(f::Function, args::AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [17] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:324
   [18] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [19] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:335 [inlined]
   [20] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [21] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:335 [inlined]
   [22] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [23] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [24] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [25] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:30 [inlined]
   [26] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [27] top-level scope
      @ none:6
   [28] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [29] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
   [30] _start()
      @ Base ./client.jl:558
with different input and output: Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:339
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [3] (::Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [4] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##reparameterize#1", Bool, Type{Float32}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [5] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [6] reparameterize
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:157 [inlined]
    [7] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.reparameterize), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:log,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{log::Bool}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##_#5", Bool, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{μ::Matrix{Float32}, σ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [9] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{μ::Matrix{Float32}, σ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [10] #loss#7
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:438 [inlined]
   [11] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [12] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:406 [inlined]
   [13] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [14] (::Zygote.Pullback{Tuple{var"#47#48"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [15] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#47#48"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [16] gradient(f::Function, args::AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [17] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:324
   [18] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [19] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:340 [inlined]
   [20] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [21] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:340 [inlined]
   [22] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [23] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [24] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [25] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:30 [inlined]
   [26] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [27] top-level scope
      @ none:6
   [28] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [29] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
   [30] _start()
      @ Base ./client.jl:558
with same input as output: Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:334
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [3] (::Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [4] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##reparameterize#1", Bool, Type{Float32}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [5] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [6] reparameterize
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:157 [inlined]
    [7] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.reparameterize), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:log,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{log::Bool}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##_#5", Bool, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [9] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [10] #loss#6
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:344 [inlined]
   [11] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [12] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:313 [inlined]
   [13] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [14] (::Zygote.Pullback{Tuple{var"#45#46"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#45#46"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [15] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#45#46"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#45#46"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [16] gradient(f::Function, args::AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [17] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:324
   [18] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [19] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:335 [inlined]
   [20] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [21] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:335 [inlined]
   [22] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [23] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [24] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [25] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:30 [inlined]
   [26] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [27] top-level scope
      @ none:6
   [28] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [29] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
   [30] _start()
      @ Base ./client.jl:558
with different input and output: Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:339
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [3] (::Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [4] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##reparameterize#1", Bool, Type{Float32}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [5] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [6] reparameterize
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:157 [inlined]
    [7] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.reparameterize), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:log,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{log::Bool}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##_#5", Bool, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [9] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [10] #loss#7
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:438 [inlined]
   [11] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [12] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:406 [inlined]
   [13] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [14] (::Zygote.Pullback{Tuple{var"#47#48"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [15] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#47#48"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [16] gradient(f::Function, args::AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [17] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:324
   [18] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [19] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:340 [inlined]
   [20] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [21] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:340 [inlined]
   [22] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [23] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [24] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [25] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:30 [inlined]
   [26] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [27] top-level scope
      @ none:6
   [28] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [29] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
   [30] _start()
      @ Base ./client.jl:558
with same input as output: Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:334
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [3] (::Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [4] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##reparameterize#1", Bool, Type{Float32}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [5] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [6] reparameterize
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:157 [inlined]
    [7] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.reparameterize), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:log,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{log::Bool}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##_#5", Bool, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{p::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [9] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{p::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [10] #loss#6
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:344 [inlined]
   [11] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}}, Matrix{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [12] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:313 [inlined]
   [13] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}}, Matrix{Float32}}, Any}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [14] (::Zygote.Pullback{Tuple{var"#45#46"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#45#46"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}}, Matrix{Float32}}, Any}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [15] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#45#46"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#45#46"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}}, Matrix{Float32}}, Any}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [16] gradient(f::Function, args::AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [17] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:324
   [18] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [19] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:335 [inlined]
   [20] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [21] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:335 [inlined]
   [22] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [23] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [24] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [25] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:30 [inlined]
   [26] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [27] top-level scope
      @ none:6
   [28] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [29] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
   [30] _start()
      @ Base ./client.jl:558
with different input and output: Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:339
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [3] (::Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [4] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##reparameterize#1", Bool, Type{Float32}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [5] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [6] reparameterize
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:157 [inlined]
    [7] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.reparameterize), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:log,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{log::Bool}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##_#5", Bool, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{p::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [9] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{p::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [10] #loss#7
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:438 [inlined]
   [11] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [12] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:406 [inlined]
   [13] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [14] (::Zygote.Pullback{Tuple{var"#47#48"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [15] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#47#48"{Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#7", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}}, Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_matrix, Zygote.Context{false}, var"#47#48"{Matrix{Float32}}, Matrix{Float32}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [16] gradient(f::Function, args::AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, BernoulliDecoder{…}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [17] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:324
   [18] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [19] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:340 [inlined]
   [20] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [21] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:340 [inlined]
   [22] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [23] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [24] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [25] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:30 [inlined]
   [26] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [27] top-level scope
      @ none:6
   [28] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [29] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
   [30] _start()
      @ Base ./client.jl:558
CPU | without regularization: Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:363
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [3] (::Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [4] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##reparameterize#1", Bool, Type{Float32}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [5] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [6] reparameterize
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:157 [inlined]
    [7] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.reparameterize), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:log,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{log::Bool}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##_#5", Bool, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [9] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Nothing, decoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [10] #loss#6
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:344 [inlined]
   [11] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, Matrix{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [12] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:313 [inlined]
   [13] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##loss#6", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, @NamedTuple{}, Float32, typeof(AutoEncoderToolkit.VAEs.loss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, Matrix{Float32}}, Any}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [14] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"#10#11"{typeof(AutoEncoderToolkit.VAEs.loss), @NamedTuple{}, Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [15] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"#10#11"{typeof(AutoEncoderToolkit.VAEs.loss), @NamedTuple{}, Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, Any}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [16] withgradient(f::Function, args::AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:213
   [17] train!(vae::AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}, x::Matrix{Float32}, opt::@NamedTuple{encoder::@NamedTuple{encoder::@NamedTuple{layers::Tuple{@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, @NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}}}, μ::@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, logσ::@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}}, decoder::@NamedTuple{decoder::@NamedTuple{layers::Tuple{@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, @NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}}}, μ::@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, logσ::@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}}}; loss_function::typeof(AutoEncoderToolkit.VAEs.loss), loss_kwargs::@NamedTuple{}, verbose::Bool, loss_return::Bool)
      @ AutoEncoderToolkit.VAEs ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:484
   [18] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:384 [inlined]
   [19] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [20] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:365 [inlined]
   [21] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [22] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:357
   [23] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [24] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [25] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [26] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:30 [inlined]
   [27] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [28] top-level scope
      @ none:6
   [29] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [30] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
   [31] _start()
      @ Base ./client.jl:558

All tests passed!


Testing MMDVAEs.jl module...

with same input and output: Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/mmdvae.jl:244
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [3] (::Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [4] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##reparameterize#1", Bool, Type{Float32}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [5] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [6] reparameterize
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:157 [inlined]
    [7] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.reparameterize), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:log,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{log::Bool}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##_#5", Bool, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [9] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [10] #_#1
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/mmdvae.jl:114 [inlined]
   [11] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.MMDVAEs.var"##_#1", Bool, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:latent,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{latent::Bool}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any}}})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [12] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [13] getindex
      @ ./tuple.jl:31 [inlined]
   [14] gradindex
      @ ~/.julia/packages/Zygote/TWpme/src/compiler/reverse.jl:12 [inlined]
   [15] #loss#5
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/mmdvae.jl:369 [inlined]
   [16] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.MMDVAEs.var"##loss#5", Float32, Float32, Int64, typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel), @NamedTuple{}, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), typeof(AutoEncoderToolkit.MMDVAEs.loss), AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Tuple{Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32, Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Any}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, @NamedTuple{μ::Matrix{Float32}}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.MMDVAEs.var"#7#8"{Int64, Matrix{Float32}, Matrix{Float32}}, Nothing, false}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##1"{Int64, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:z, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:latent,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{latent::Bool}, Nothing, true}}}}, Zygote.ZBack{ChainRules.MergePullback{@NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, @NamedTuple{}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Matrix{Float32}}, Any}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), Matrix{Float32}, Matrix{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.var"##decoder_loglikelihood#24", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), Matrix{Float32}, Matrix{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{typeof(vec), Matrix{Float32}}, Tuple{Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.Pullback{Tuple{typeof(ndims), Matrix{Float32}}, Tuple{Zygote.var"#188#189"{Zygote.var"#186#187"}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32}}}, Zygote.var"#1535#1536"{Zygote.var"#1533#1534"}, Zygote.ZBack{Zygote.var"#convert_pullback#convert_pullback##0"}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{UnitRange{Int64}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::UnitRange{Int64}}, Nothing, true}}}}, Zygote.ZBack{ChainRules.var"#log_pullback#log_pullback##0"{Float32, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#log_pullback#log_pullback##0"{Float32, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Irrational{:π}}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#slash_pullback_scalar#rrule##224"{Matrix{Float32}, Float32}}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.ZBack{ChainRules.var"#getindex_pullback#rrule##249"{Matrix{Float32}, Tuple{EllipsisNotation.Ellipsis, Int64}, Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Float32}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.ZBack{Zygote.var"#literal_pow_pullback#literal_pow_pullback##0"{2, Float32}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{ChainRules.var"#times_pullback3#rrule##58"{Float32, Float32, Float32}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(eltype), Matrix{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:z, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:encoder, Zygote.Context{false}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, JointGaussianLogEncoder{…}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:kernel,)}}, Tuple{typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, Nothing, true}}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, typeof(AutoEncoderToolkit.MMDVAEs.mmd_div), Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.ZBack{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Colon, Vector{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.encoder_kl), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcasted), typeof(exp), Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcastable), Matrix{Float32}}, Tuple{}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing, Nothing, Nothing}, Tuple{}}, Zygote.var"#1949#1950"{Zygote.var"#bc_fwd_back#_broadcast_forward##2"{Matrix{ForwardDiff.Dual{Nothing, Float32, 1}}, Tuple{Matrix{Float32}}, Val{1}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#890#891"{Zygote.var"#map_back#∇map##2"{typeof(Base.Broadcast.broadcastable), 1, Tuple{Tuple{}}, Tuple{Val{0}}, Tuple{}}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing}, Tuple{}}, Zygote.var"#combine_styles_pullback#combine_styles_pullback##0"{Tuple{Nothing, Nothing}}}}}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(vec), Matrix{Float32}}, Tuple{Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Matrix{Float32}}}, Zygote.var"#1707#1708"{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Matrix{Float32}, Matrix{Float32}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.var"#1707#1708"{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Float32}}, Zygote.var"#162#163"{Zygote.var"#160#161"}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, SimpleGaussianDecoder{…}}}, Zygote.ZBack{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Colon, Vector{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:encoder, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [17] (::Zygote.Pullback{Tuple{var"#53#54"{Matrix{Float32}}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, var"#53#54"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.MMDVAEs.loss), AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.MMDVAEs.var"##loss#5", Float32, Float32, Int64, typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel), @NamedTuple{}, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), typeof(AutoEncoderToolkit.MMDVAEs.loss), AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Tuple{Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32, Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Any}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, @NamedTuple{μ::Matrix{Float32}}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.MMDVAEs.var"#7#8"{Int64, Matrix{Float32}, Matrix{Float32}}, Nothing, false}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##1"{Int64, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:z, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:latent,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{latent::Bool}, Nothing, true}}}}, Zygote.ZBack{ChainRules.MergePullback{@NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, @NamedTuple{}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Matrix{Float32}}, Any}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), Matrix{Float32}, Matrix{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.var"##decoder_loglikelihood#24", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), Matrix{Float32}, Matrix{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{typeof(vec), Matrix{Float32}}, Tuple{Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.Pullback{Tuple{typeof(ndims), Matrix{Float32}}, Tuple{Zygote.var"#188#189"{Zygote.var"#186#187"}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32}}}, Zygote.var"#1535#1536"{Zygote.var"#1533#1534"}, Zygote.ZBack{Zygote.var"#convert_pullback#convert_pullback##0"}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{UnitRange{Int64}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::UnitRange{Int64}}, Nothing, true}}}}, Zygote.ZBack{ChainRules.var"#log_pullback#log_pullback##0"{Float32, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#log_pullback#log_pullback##0"{Float32, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Irrational{:π}}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#slash_pullback_scalar#rrule##224"{Matrix{Float32}, Float32}}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.ZBack{ChainRules.var"#getindex_pullback#rrule##249"{Matrix{Float32}, Tuple{EllipsisNotation.Ellipsis, Int64}, Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Float32}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.ZBack{Zygote.var"#literal_pow_pullback#literal_pow_pullback##0"{2, Float32}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{ChainRules.var"#times_pullback3#rrule##58"{Float32, Float32, Float32}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(eltype), Matrix{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:z, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:encoder, Zygote.Context{false}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, JointGaussianLogEncoder{…}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:kernel,)}}, Tuple{typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, Nothing, true}}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, typeof(AutoEncoderToolkit.MMDVAEs.mmd_div), Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.ZBack{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Colon, Vector{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.encoder_kl), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcasted), typeof(exp), Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcastable), Matrix{Float32}}, Tuple{}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing, Nothing, Nothing}, Tuple{}}, Zygote.var"#1949#1950"{Zygote.var"#bc_fwd_back#_broadcast_forward##2"{Matrix{ForwardDiff.Dual{Nothing, Float32, 1}}, Tuple{Matrix{Float32}}, Val{1}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#890#891"{Zygote.var"#map_back#∇map##2"{typeof(Base.Broadcast.broadcastable), 1, Tuple{Tuple{}}, Tuple{Val{0}}, Tuple{}}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing}, Tuple{}}, Zygote.var"#combine_styles_pullback#combine_styles_pullback##0"{Tuple{Nothing, Nothing}}}}}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(vec), Matrix{Float32}}, Tuple{Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Matrix{Float32}}}, Zygote.var"#1707#1708"{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Matrix{Float32}, Matrix{Float32}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.var"#1707#1708"{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Float32}}, Zygote.var"#162#163"{Zygote.var"#160#161"}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, SimpleGaussianDecoder{…}}}, Zygote.ZBack{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Colon, Vector{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:encoder, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [18] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#53#54"{Matrix{Float32}}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, var"#53#54"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.MMDVAEs.loss), AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.MMDVAEs.var"##loss#5", Float32, Float32, Int64, typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel), @NamedTuple{}, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), typeof(AutoEncoderToolkit.MMDVAEs.loss), AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Tuple{Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32, Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Any}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, @NamedTuple{μ::Matrix{Float32}}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.MMDVAEs.var"#7#8"{Int64, Matrix{Float32}, Matrix{Float32}}, Nothing, false}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##1"{Int64, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:z, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:latent,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{latent::Bool}, Nothing, true}}}}, Zygote.ZBack{ChainRules.MergePullback{@NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, @NamedTuple{}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Matrix{Float32}}, Any}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), Matrix{Float32}, Matrix{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.var"##decoder_loglikelihood#24", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), Matrix{Float32}, Matrix{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{typeof(vec), Matrix{Float32}}, Tuple{Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.Pullback{Tuple{typeof(ndims), Matrix{Float32}}, Tuple{Zygote.var"#188#189"{Zygote.var"#186#187"}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32}}}, Zygote.var"#1535#1536"{Zygote.var"#1533#1534"}, Zygote.ZBack{Zygote.var"#convert_pullback#convert_pullback##0"}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{UnitRange{Int64}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::UnitRange{Int64}}, Nothing, true}}}}, Zygote.ZBack{ChainRules.var"#log_pullback#log_pullback##0"{Float32, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#log_pullback#log_pullback##0"{Float32, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Irrational{:π}}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#slash_pullback_scalar#rrule##224"{Matrix{Float32}, Float32}}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.ZBack{ChainRules.var"#getindex_pullback#rrule##249"{Matrix{Float32}, Tuple{EllipsisNotation.Ellipsis, Int64}, Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Float32}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.ZBack{Zygote.var"#literal_pow_pullback#literal_pow_pullback##0"{2, Float32}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{ChainRules.var"#times_pullback3#rrule##58"{Float32, Float32, Float32}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(eltype), Matrix{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:z, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:encoder, Zygote.Context{false}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, JointGaussianLogEncoder{…}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:kernel,)}}, Tuple{typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, Nothing, true}}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, typeof(AutoEncoderToolkit.MMDVAEs.mmd_div), Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.ZBack{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Colon, Vector{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.encoder_kl), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcasted), typeof(exp), Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcastable), Matrix{Float32}}, Tuple{}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing, Nothing, Nothing}, Tuple{}}, Zygote.var"#1949#1950"{Zygote.var"#bc_fwd_back#_broadcast_forward##2"{Matrix{ForwardDiff.Dual{Nothing, Float32, 1}}, Tuple{Matrix{Float32}}, Val{1}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#890#891"{Zygote.var"#map_back#∇map##2"{typeof(Base.Broadcast.broadcastable), 1, Tuple{Tuple{}}, Tuple{Val{0}}, Tuple{}}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing}, Tuple{}}, Zygote.var"#combine_styles_pullback#combine_styles_pullback##0"{Tuple{Nothing, Nothing}}}}}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(vec), Matrix{Float32}}, Tuple{Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Matrix{Float32}}}, Zygote.var"#1707#1708"{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Matrix{Float32}, Matrix{Float32}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.var"#1707#1708"{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Float32}}, Zygote.var"#162#163"{Zygote.var"#160#161"}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, SimpleGaussianDecoder{…}}}, Zygote.ZBack{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Colon, Vector{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:encoder, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}}}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [19] gradient(f::Function, args::AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [20] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/mmdvae.jl:215
   [21] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [22] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/mmdvae.jl:242 [inlined]
   [23] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [24] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/mmdvae.jl:245 [inlined]
   [25] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [26] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/mmdvae.jl:245 [inlined]
   [27] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [28] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [29] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [30] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:33 [inlined]
   [31] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [32] top-level scope
      @ none:6
   [33] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [34] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
   [35] _start()
      @ Base ./client.jl:558
with different input and output: Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/mmdvae.jl:249
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [3] (::Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [4] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##reparameterize#1", Bool, Type{Float32}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [5] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [6] reparameterize
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:157 [inlined]
    [7] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.reparameterize), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:log,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{log::Bool}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##_#5", Bool, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [9] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [10] #_#1
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/mmdvae.jl:114 [inlined]
   [11] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.MMDVAEs.var"##_#1", Bool, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:latent,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{latent::Bool}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any}}})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [12] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [13] getindex
      @ ./tuple.jl:31 [inlined]
   [14] gradindex
      @ ~/.julia/packages/Zygote/TWpme/src/compiler/reverse.jl:12 [inlined]
   [15] #loss#9
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/mmdvae.jl:466 [inlined]
   [16] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.MMDVAEs.var"##loss#9", Float32, Float32, Int64, typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel), @NamedTuple{}, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), typeof(AutoEncoderToolkit.MMDVAEs.loss), AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.encoder_kl), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcasted), typeof(exp), Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcastable), Matrix{Float32}}, Tuple{}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing, Nothing, Nothing}, Tuple{}}, Zygote.var"#1949#1950"{Zygote.var"#bc_fwd_back#_broadcast_forward##2"{Matrix{ForwardDiff.Dual{Nothing, Float32, 1}}, Tuple{Matrix{Float32}}, Val{1}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#890#891"{Zygote.var"#map_back#∇map##2"{typeof(Base.Broadcast.broadcastable), 1, Tuple{Tuple{}}, Tuple{Val{0}}, Tuple{}}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing}, Tuple{}}, Zygote.var"#combine_styles_pullback#combine_styles_pullback##0"{Tuple{Nothing, Nothing}}}}}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(vec), Matrix{Float32}}, Tuple{Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Matrix{Float32}}}, Zygote.var"#1707#1708"{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Matrix{Float32}, Matrix{Float32}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.var"#1707#1708"{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Float32}}, Zygote.var"#162#163"{Zygote.var"#160#161"}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:encoder, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}}, Zygote.ZBack{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Colon, Vector{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Zygote.ZBack{ChainRules.MergePullback{@NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, @NamedTuple{}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:kernel,)}}, Tuple{typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, Nothing, true}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, SimpleGaussianDecoder{…}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:z, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Colon, Vector{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, typeof(AutoEncoderToolkit.MMDVAEs.mmd_div), Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:encoder, Zygote.Context{false}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, JointGaussianLogEncoder{…}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:latent,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{latent::Bool}, Nothing, true}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##1"{Int64, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32, Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), Matrix{Float32}, Matrix{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.var"##decoder_loglikelihood#24", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), Matrix{Float32}, Matrix{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{typeof(vec), Matrix{Float32}}, Tuple{Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.Pullback{Tuple{typeof(ndims), Matrix{Float32}}, Tuple{Zygote.var"#188#189"{Zygote.var"#186#187"}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32}}}, Zygote.var"#1535#1536"{Zygote.var"#1533#1534"}, Zygote.ZBack{Zygote.var"#convert_pullback#convert_pullback##0"}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{UnitRange{Int64}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::UnitRange{Int64}}, Nothing, true}}}}, Zygote.ZBack{ChainRules.var"#log_pullback#log_pullback##0"{Float32, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#log_pullback#log_pullback##0"{Float32, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Irrational{:π}}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#slash_pullback_scalar#rrule##224"{Matrix{Float32}, Float32}}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.ZBack{ChainRules.var"#getindex_pullback#rrule##249"{Matrix{Float32}, Tuple{EllipsisNotation.Ellipsis, Int64}, Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Float32}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.ZBack{Zygote.var"#literal_pow_pullback#literal_pow_pullback##0"{2, Float32}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{ChainRules.var"#times_pullback3#rrule##58"{Float32, Float32, Float32}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(eltype), Matrix{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:z, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, Matrix{Float32}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.MMDVAEs.var"#11#12"{Int64, Matrix{Float32}, Matrix{Float32}}, Nothing, false}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, @NamedTuple{μ::Matrix{Float32}}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Matrix{Float32}}, Any}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [17] (::Zygote.Pullback{Tuple{var"#55#56"{Matrix{Float32}, Matrix{Float32}}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, var"#55#56"{Matrix{Float32}, Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.MMDVAEs.loss), AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.MMDVAEs.var"##loss#9", Float32, Float32, Int64, typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel), @NamedTuple{}, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), typeof(AutoEncoderToolkit.MMDVAEs.loss), AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.encoder_kl), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcasted), typeof(exp), Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcastable), Matrix{Float32}}, Tuple{}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing, Nothing, Nothing}, Tuple{}}, Zygote.var"#1949#1950"{Zygote.var"#bc_fwd_back#_broadcast_forward##2"{Matrix{ForwardDiff.Dual{Nothing, Float32, 1}}, Tuple{Matrix{Float32}}, Val{1}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#890#891"{Zygote.var"#map_back#∇map##2"{typeof(Base.Broadcast.broadcastable), 1, Tuple{Tuple{}}, Tuple{Val{0}}, Tuple{}}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing}, Tuple{}}, Zygote.var"#combine_styles_pullback#combine_styles_pullback##0"{Tuple{Nothing, Nothing}}}}}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(vec), Matrix{Float32}}, Tuple{Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Matrix{Float32}}}, Zygote.var"#1707#1708"{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Matrix{Float32}, Matrix{Float32}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.var"#1707#1708"{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Float32}}, Zygote.var"#162#163"{Zygote.var"#160#161"}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:encoder, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}}, Zygote.ZBack{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Colon, Vector{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Zygote.ZBack{ChainRules.MergePullback{@NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, @NamedTuple{}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:kernel,)}}, Tuple{typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, Nothing, true}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, SimpleGaussianDecoder{…}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:z, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Colon, Vector{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, typeof(AutoEncoderToolkit.MMDVAEs.mmd_div), Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:encoder, Zygote.Context{false}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, JointGaussianLogEncoder{…}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:latent,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{latent::Bool}, Nothing, true}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##1"{Int64, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32, Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), Matrix{Float32}, Matrix{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.var"##decoder_loglikelihood#24", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), Matrix{Float32}, Matrix{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{typeof(vec), Matrix{Float32}}, Tuple{Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.Pullback{Tuple{typeof(ndims), Matrix{Float32}}, Tuple{Zygote.var"#188#189"{Zygote.var"#186#187"}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32}}}, Zygote.var"#1535#1536"{Zygote.var"#1533#1534"}, Zygote.ZBack{Zygote.var"#convert_pullback#convert_pullback##0"}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{UnitRange{Int64}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::UnitRange{Int64}}, Nothing, true}}}}, Zygote.ZBack{ChainRules.var"#log_pullback#log_pullback##0"{Float32, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#log_pullback#log_pullback##0"{Float32, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Irrational{:π}}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#slash_pullback_scalar#rrule##224"{Matrix{Float32}, Float32}}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.ZBack{ChainRules.var"#getindex_pullback#rrule##249"{Matrix{Float32}, Tuple{EllipsisNotation.Ellipsis, Int64}, Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Float32}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.ZBack{Zygote.var"#literal_pow_pullback#literal_pow_pullback##0"{2, Float32}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{ChainRules.var"#times_pullback3#rrule##58"{Float32, Float32, Float32}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(eltype), Matrix{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:z, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, Matrix{Float32}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.MMDVAEs.var"#11#12"{Int64, Matrix{Float32}, Matrix{Float32}}, Nothing, false}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, @NamedTuple{μ::Matrix{Float32}}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Matrix{Float32}}, Any}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_out, Zygote.Context{false}, var"#55#56"{Matrix{Float32}, Matrix{Float32}}, Matrix{Float32}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [18] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#55#56"{Matrix{Float32}, Matrix{Float32}}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, var"#55#56"{Matrix{Float32}, Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.MMDVAEs.loss), AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.MMDVAEs.var"##loss#9", Float32, Float32, Int64, typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel), @NamedTuple{}, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), typeof(AutoEncoderToolkit.MMDVAEs.loss), AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.encoder_kl), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcasted), typeof(exp), Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcastable), Matrix{Float32}}, Tuple{}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing, Nothing, Nothing}, Tuple{}}, Zygote.var"#1949#1950"{Zygote.var"#bc_fwd_back#_broadcast_forward##2"{Matrix{ForwardDiff.Dual{Nothing, Float32, 1}}, Tuple{Matrix{Float32}}, Val{1}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#890#891"{Zygote.var"#map_back#∇map##2"{typeof(Base.Broadcast.broadcastable), 1, Tuple{Tuple{}}, Tuple{Val{0}}, Tuple{}}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing}, Tuple{}}, Zygote.var"#combine_styles_pullback#combine_styles_pullback##0"{Tuple{Nothing, Nothing}}}}}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(vec), Matrix{Float32}}, Tuple{Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Matrix{Float32}}}, Zygote.var"#1707#1708"{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Matrix{Float32}, Matrix{Float32}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.var"#1707#1708"{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Float32}}, Zygote.var"#162#163"{Zygote.var"#160#161"}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:encoder, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}}, Zygote.ZBack{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Colon, Vector{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Zygote.ZBack{ChainRules.MergePullback{@NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, @NamedTuple{}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:kernel,)}}, Tuple{typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, Nothing, true}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, SimpleGaussianDecoder{…}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:z, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Colon, Vector{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, typeof(AutoEncoderToolkit.MMDVAEs.mmd_div), Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:encoder, Zygote.Context{false}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, JointGaussianLogEncoder{…}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:latent,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{latent::Bool}, Nothing, true}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##1"{Int64, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32, Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), Matrix{Float32}, Matrix{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.var"##decoder_loglikelihood#24", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), Matrix{Float32}, Matrix{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{typeof(vec), Matrix{Float32}}, Tuple{Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.Pullback{Tuple{typeof(ndims), Matrix{Float32}}, Tuple{Zygote.var"#188#189"{Zygote.var"#186#187"}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32}}}, Zygote.var"#1535#1536"{Zygote.var"#1533#1534"}, Zygote.ZBack{Zygote.var"#convert_pullback#convert_pullback##0"}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{UnitRange{Int64}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::UnitRange{Int64}}, Nothing, true}}}}, Zygote.ZBack{ChainRules.var"#log_pullback#log_pullback##0"{Float32, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#log_pullback#log_pullback##0"{Float32, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Irrational{:π}}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#slash_pullback_scalar#rrule##224"{Matrix{Float32}, Float32}}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.ZBack{ChainRules.var"#getindex_pullback#rrule##249"{Matrix{Float32}, Tuple{EllipsisNotation.Ellipsis, Int64}, Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Float32}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.ZBack{Zygote.var"#literal_pow_pullback#literal_pow_pullback##0"{2, Float32}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{ChainRules.var"#times_pullback3#rrule##58"{Float32, Float32, Float32}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(eltype), Matrix{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:z, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, Matrix{Float32}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.MMDVAEs.var"#11#12"{Int64, Matrix{Float32}, Matrix{Float32}}, Nothing, false}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, @NamedTuple{μ::Matrix{Float32}}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Matrix{Float32}}, Any}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_out, Zygote.Context{false}, var"#55#56"{Matrix{Float32}, Matrix{Float32}}, Matrix{Float32}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [19] gradient(f::Function, args::AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [20] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/mmdvae.jl:215
   [21] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [22] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/mmdvae.jl:242 [inlined]
   [23] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [24] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/mmdvae.jl:250 [inlined]
   [25] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [26] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/mmdvae.jl:251 [inlined]
   [27] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [28] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [29] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [30] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:33 [inlined]
   [31] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [32] top-level scope
      @ none:6
   [33] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [34] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
   [35] _start()
      @ Base ./client.jl:558
with same input and output: Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/mmdvae.jl:323
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [3] (::Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [4] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##reparameterize#1", Bool, Type{Float32}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [5] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [6] reparameterize
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:157 [inlined]
    [7] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.reparameterize), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:log,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{log::Bool}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##_#5", Bool, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [9] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [10] #_#1
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/mmdvae.jl:114 [inlined]
   [11] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.MMDVAEs.var"##_#1", Bool, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:latent,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{latent::Bool}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any}}})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [12] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [13] getindex
      @ ./tuple.jl:31 [inlined]
   [14] gradindex
      @ ~/.julia/packages/Zygote/TWpme/src/compiler/reverse.jl:12 [inlined]
   [15] #loss#5
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/mmdvae.jl:369 [inlined]
   [16] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.MMDVAEs.var"##loss#5", Float32, Float32, Int64, typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel), @NamedTuple{}, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), typeof(AutoEncoderToolkit.MMDVAEs.loss), AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Tuple{Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32, Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Any}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, @NamedTuple{μ::Matrix{Float32}}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.MMDVAEs.var"#7#8"{Int64, Matrix{Float32}, Matrix{Float32}}, Nothing, false}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##1"{Int64, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:z, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:latent,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{latent::Bool}, Nothing, true}}}}, Zygote.ZBack{ChainRules.MergePullback{@NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, @NamedTuple{}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Matrix{Float32}}, Any}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), Matrix{Float32}, Matrix{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.var"##decoder_loglikelihood#24", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), Matrix{Float32}, Matrix{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{typeof(vec), Matrix{Float32}}, Tuple{Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.Pullback{Tuple{typeof(ndims), Matrix{Float32}}, Tuple{Zygote.var"#188#189"{Zygote.var"#186#187"}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32}}}, Zygote.var"#1535#1536"{Zygote.var"#1533#1534"}, Zygote.ZBack{Zygote.var"#convert_pullback#convert_pullback##0"}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{UnitRange{Int64}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::UnitRange{Int64}}, Nothing, true}}}}, Zygote.ZBack{ChainRules.var"#log_pullback#log_pullback##0"{Float32, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#log_pullback#log_pullback##0"{Float32, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Irrational{:π}}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#slash_pullback_scalar#rrule##224"{Matrix{Float32}, Float32}}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.ZBack{ChainRules.var"#getindex_pullback#rrule##249"{Matrix{Float32}, Tuple{EllipsisNotation.Ellipsis, Int64}, Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Float32}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.ZBack{Zygote.var"#literal_pow_pullback#literal_pow_pullback##0"{2, Float32}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{ChainRules.var"#times_pullback3#rrule##58"{Float32, Float32, Float32}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(eltype), Matrix{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:z, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:encoder, Zygote.Context{false}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, JointGaussianLogEncoder{…}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:kernel,)}}, Tuple{typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, Nothing, true}}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, typeof(AutoEncoderToolkit.MMDVAEs.mmd_div), Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.ZBack{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Colon, Vector{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.encoder_kl), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcasted), typeof(exp), Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcastable), Matrix{Float32}}, Tuple{}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing, Nothing, Nothing}, Tuple{}}, Zygote.var"#1949#1950"{Zygote.var"#bc_fwd_back#_broadcast_forward##2"{Matrix{ForwardDiff.Dual{Nothing, Float32, 1}}, Tuple{Matrix{Float32}}, Val{1}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#890#891"{Zygote.var"#map_back#∇map##2"{typeof(Base.Broadcast.broadcastable), 1, Tuple{Tuple{}}, Tuple{Val{0}}, Tuple{}}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing}, Tuple{}}, Zygote.var"#combine_styles_pullback#combine_styles_pullback##0"{Tuple{Nothing, Nothing}}}}}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(vec), Matrix{Float32}}, Tuple{Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Matrix{Float32}}}, Zygote.var"#1707#1708"{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Matrix{Float32}, Matrix{Float32}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.var"#1707#1708"{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Float32}}, Zygote.var"#162#163"{Zygote.var"#160#161"}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, SimpleGaussianDecoder{…}}}, Zygote.ZBack{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Colon, Vector{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:encoder, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [17] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/mmdvae.jl:328 [inlined]
   [18] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.MMDVAEs.loss), AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.MMDVAEs.var"##loss#5", Float32, Float32, Int64, typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel), @NamedTuple{}, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), typeof(AutoEncoderToolkit.MMDVAEs.loss), AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Tuple{Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32, Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Any}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, @NamedTuple{μ::Matrix{Float32}}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.MMDVAEs.var"#7#8"{Int64, Matrix{Float32}, Matrix{Float32}}, Nothing, false}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##1"{Int64, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:z, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:latent,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{latent::Bool}, Nothing, true}}}}, Zygote.ZBack{ChainRules.MergePullback{@NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, @NamedTuple{}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Matrix{Float32}}, Any}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), Matrix{Float32}, Matrix{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.var"##decoder_loglikelihood#24", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), Matrix{Float32}, Matrix{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{typeof(vec), Matrix{Float32}}, Tuple{Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.Pullback{Tuple{typeof(ndims), Matrix{Float32}}, Tuple{Zygote.var"#188#189"{Zygote.var"#186#187"}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32}}}, Zygote.var"#1535#1536"{Zygote.var"#1533#1534"}, Zygote.ZBack{Zygote.var"#convert_pullback#convert_pullback##0"}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{UnitRange{Int64}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::UnitRange{Int64}}, Nothing, true}}}}, Zygote.ZBack{ChainRules.var"#log_pullback#log_pullback##0"{Float32, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#log_pullback#log_pullback##0"{Float32, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Irrational{:π}}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#slash_pullback_scalar#rrule##224"{Matrix{Float32}, Float32}}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.ZBack{ChainRules.var"#getindex_pullback#rrule##249"{Matrix{Float32}, Tuple{EllipsisNotation.Ellipsis, Int64}, Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Float32}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.ZBack{Zygote.var"#literal_pow_pullback#literal_pow_pullback##0"{2, Float32}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{ChainRules.var"#times_pullback3#rrule##58"{Float32, Float32, Float32}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(eltype), Matrix{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:z, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:encoder, Zygote.Context{false}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, JointGaussianLogEncoder{…}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:kernel,)}}, Tuple{typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, Nothing, true}}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, typeof(AutoEncoderToolkit.MMDVAEs.mmd_div), Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.ZBack{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Colon, Vector{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.encoder_kl), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcasted), typeof(exp), Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcastable), Matrix{Float32}}, Tuple{}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing, Nothing, Nothing}, Tuple{}}, Zygote.var"#1949#1950"{Zygote.var"#bc_fwd_back#_broadcast_forward##2"{Matrix{ForwardDiff.Dual{Nothing, Float32, 1}}, Tuple{Matrix{Float32}}, Val{1}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#890#891"{Zygote.var"#map_back#∇map##2"{typeof(Base.Broadcast.broadcastable), 1, Tuple{Tuple{}}, Tuple{Val{0}}, Tuple{}}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing}, Tuple{}}, Zygote.var"#combine_styles_pullback#combine_styles_pullback##0"{Tuple{Nothing, Nothing}}}}}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(vec), Matrix{Float32}}, Tuple{Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Matrix{Float32}}}, Zygote.var"#1707#1708"{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Matrix{Float32}, Matrix{Float32}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.var"#1707#1708"{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Float32}}, Zygote.var"#162#163"{Zygote.var"#160#161"}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, SimpleGaussianDecoder{…}}}, Zygote.ZBack{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Colon, Vector{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:encoder, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [19] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.MMDVAEs.var"#15#16"{typeof(AutoEncoderToolkit.MMDVAEs.loss), @NamedTuple{}, Matrix{Float32}}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [20] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{AutoEncoderToolkit.MMDVAEs.var"#15#16"{typeof(AutoEncoderToolkit.MMDVAEs.loss), @NamedTuple{}, Matrix{Float32}}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Any}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [21] withgradient(f::Function, args::AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:213
   [22] train!(mmdvae::AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, x::Matrix{Float32}, opt::@NamedTuple{vae::@NamedTuple{encoder::@NamedTuple{encoder::@NamedTuple{layers::Tuple{@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, @NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}}}, μ::@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, logσ::@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}}, decoder::@NamedTuple{decoder::@NamedTuple{layers::Tuple{@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, @NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, @NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}}}}}}; loss_function::typeof(AutoEncoderToolkit.MMDVAEs.loss), loss_kwargs::@NamedTuple{}, verbose::Bool, loss_return::Bool)
      @ AutoEncoderToolkit.MMDVAEs ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/mmdvae.jl:513
   [23] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/mmdvae.jl:291
   [24] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [25] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/mmdvae.jl:318 [inlined]
   [26] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [27] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/mmdvae.jl:324 [inlined]
   [28] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [29] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/mmdvae.jl:324 [inlined]
   [30] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [31] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [32] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [33] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:33 [inlined]
   [34] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [35] top-level scope
      @ none:6
   [36] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [37] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
   [38] _start()
      @ Base ./client.jl:558
with different input and output: Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/mmdvae.jl:328
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [3] (::Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [4] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##reparameterize#1", Bool, Type{Float32}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [5] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [6] reparameterize
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:157 [inlined]
    [7] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.reparameterize), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:log,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{log::Bool}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##_#5", Bool, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [9] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [10] #_#1
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/mmdvae.jl:114 [inlined]
   [11] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.MMDVAEs.var"##_#1", Bool, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:latent,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{latent::Bool}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any}}})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [12] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [13] getindex
      @ ./tuple.jl:31 [inlined]
   [14] gradindex
      @ ~/.julia/packages/Zygote/TWpme/src/compiler/reverse.jl:12 [inlined]
   [15] #loss#9
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/mmdvae.jl:466 [inlined]
   [16] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.MMDVAEs.var"##loss#9", Float32, Float32, Int64, typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel), @NamedTuple{}, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), typeof(AutoEncoderToolkit.MMDVAEs.loss), AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.encoder_kl), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcasted), typeof(exp), Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcastable), Matrix{Float32}}, Tuple{}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing, Nothing, Nothing}, Tuple{}}, Zygote.var"#1949#1950"{Zygote.var"#bc_fwd_back#_broadcast_forward##2"{Matrix{ForwardDiff.Dual{Nothing, Float32, 1}}, Tuple{Matrix{Float32}}, Val{1}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#890#891"{Zygote.var"#map_back#∇map##2"{typeof(Base.Broadcast.broadcastable), 1, Tuple{Tuple{}}, Tuple{Val{0}}, Tuple{}}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing}, Tuple{}}, Zygote.var"#combine_styles_pullback#combine_styles_pullback##0"{Tuple{Nothing, Nothing}}}}}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(vec), Matrix{Float32}}, Tuple{Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Matrix{Float32}}}, Zygote.var"#1707#1708"{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Matrix{Float32}, Matrix{Float32}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.var"#1707#1708"{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Float32}}, Zygote.var"#162#163"{Zygote.var"#160#161"}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:encoder, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}}, Zygote.ZBack{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Colon, Vector{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Zygote.ZBack{ChainRules.MergePullback{@NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, @NamedTuple{}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:kernel,)}}, Tuple{typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, Nothing, true}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, SimpleGaussianDecoder{…}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:z, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Colon, Vector{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, typeof(AutoEncoderToolkit.MMDVAEs.mmd_div), Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:encoder, Zygote.Context{false}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, JointGaussianLogEncoder{…}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:latent,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{latent::Bool}, Nothing, true}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##1"{Int64, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32, Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), Matrix{Float32}, Matrix{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.var"##decoder_loglikelihood#24", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), Matrix{Float32}, Matrix{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{typeof(vec), Matrix{Float32}}, Tuple{Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.Pullback{Tuple{typeof(ndims), Matrix{Float32}}, Tuple{Zygote.var"#188#189"{Zygote.var"#186#187"}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32}}}, Zygote.var"#1535#1536"{Zygote.var"#1533#1534"}, Zygote.ZBack{Zygote.var"#convert_pullback#convert_pullback##0"}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{UnitRange{Int64}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::UnitRange{Int64}}, Nothing, true}}}}, Zygote.ZBack{ChainRules.var"#log_pullback#log_pullback##0"{Float32, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#log_pullback#log_pullback##0"{Float32, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Irrational{:π}}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#slash_pullback_scalar#rrule##224"{Matrix{Float32}, Float32}}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.ZBack{ChainRules.var"#getindex_pullback#rrule##249"{Matrix{Float32}, Tuple{EllipsisNotation.Ellipsis, Int64}, Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Float32}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.ZBack{Zygote.var"#literal_pow_pullback#literal_pow_pullback##0"{2, Float32}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{ChainRules.var"#times_pullback3#rrule##58"{Float32, Float32, Float32}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(eltype), Matrix{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:z, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, Matrix{Float32}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.MMDVAEs.var"#11#12"{Int64, Matrix{Float32}, Matrix{Float32}}, Nothing, false}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, @NamedTuple{μ::Matrix{Float32}}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Matrix{Float32}}, Any}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [17] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/mmdvae.jl:424 [inlined]
   [18] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.MMDVAEs.loss), AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.MMDVAEs.var"##loss#9", Float32, Float32, Int64, typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel), @NamedTuple{}, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), typeof(AutoEncoderToolkit.MMDVAEs.loss), AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.encoder_kl), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcasted), typeof(exp), Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcastable), Matrix{Float32}}, Tuple{}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing, Nothing, Nothing}, Tuple{}}, Zygote.var"#1949#1950"{Zygote.var"#bc_fwd_back#_broadcast_forward##2"{Matrix{ForwardDiff.Dual{Nothing, Float32, 1}}, Tuple{Matrix{Float32}}, Val{1}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#890#891"{Zygote.var"#map_back#∇map##2"{typeof(Base.Broadcast.broadcastable), 1, Tuple{Tuple{}}, Tuple{Val{0}}, Tuple{}}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing}, Tuple{}}, Zygote.var"#combine_styles_pullback#combine_styles_pullback##0"{Tuple{Nothing, Nothing}}}}}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(vec), Matrix{Float32}}, Tuple{Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Matrix{Float32}}}, Zygote.var"#1707#1708"{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Matrix{Float32}, Matrix{Float32}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.var"#1707#1708"{Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Float32}}, Zygote.var"#162#163"{Zygote.var"#160#161"}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:encoder, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}}, Zygote.ZBack{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Colon, Vector{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Zygote.ZBack{ChainRules.MergePullback{@NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, @NamedTuple{}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:kernel,)}}, Tuple{typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, Nothing, true}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, SimpleGaussianDecoder{…}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:z, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Colon, Vector{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:vae, Zygote.Context{false}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{kernel::typeof(AutoEncoderToolkit.MMDVAEs.gaussian_kernel)}, typeof(AutoEncoderToolkit.MMDVAEs.mmd_div), Matrix{Float32}, Matrix{Float32}}, Any}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:encoder, Zygote.Context{false}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, JointGaussianLogEncoder{…}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:latent,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{latent::Bool}, Nothing, true}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##1"{Int64, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32, Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), Matrix{Float32}, Matrix{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.var"##decoder_loglikelihood#24", Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), Matrix{Float32}, Matrix{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{typeof(vec), Matrix{Float32}}, Tuple{Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.Pullback{Tuple{typeof(ndims), Matrix{Float32}}, Tuple{Zygote.var"#188#189"{Zygote.var"#186#187"}}}, Zygote.ZBack{Zygote.var"#plus_pullback#plus_pullback##0"{Tuple{Float32, Float32}}}, Zygote.var"#1535#1536"{Zygote.var"#1533#1534"}, Zygote.ZBack{Zygote.var"#convert_pullback#convert_pullback##0"}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{UnitRange{Int64}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::UnitRange{Int64}}, Nothing, true}}}}, Zygote.ZBack{ChainRules.var"#log_pullback#log_pullback##0"{Float32, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##0"{Int64, Bool, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}, ChainRulesCore.ProjectTo{Float64, @NamedTuple{}}}}, Zygote.ZBack{ChainRules.var"#log_pullback#log_pullback##0"{Float32, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Irrational{:π}}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#slash_pullback_scalar#rrule##224"{Matrix{Float32}, Float32}}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.ZBack{ChainRules.var"#getindex_pullback#rrule##249"{Matrix{Float32}, Tuple{EllipsisNotation.Ellipsis, Int64}, Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Matrix{Float32}, Float32}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}, Zygote.ZBack{Zygote.var"#literal_pow_pullback#literal_pow_pullback##0"{2, Float32}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##147"{Float32, Matrix{Float32}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{ChainRules.var"#times_pullback3#rrule##58"{Float32, Float32, Float32}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(eltype), Matrix{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:z, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, Matrix{Float32}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.MMDVAEs.var"#11#12"{Int64, Matrix{Float32}, Matrix{Float32}}, Nothing, false}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, @NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}}, @NamedTuple{μ::Matrix{Float32}}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Matrix{Float32}}, Any}, Zygote.ZBack{ChainRules.var"#times_pullback2#rrule##57"{Float32, Float32}}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [19] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.MMDVAEs.var"#19#20"{typeof(AutoEncoderToolkit.MMDVAEs.loss), @NamedTuple{}, Matrix{Float32}, Matrix{Float32}}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [20] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{AutoEncoderToolkit.MMDVAEs.var"#19#20"{typeof(AutoEncoderToolkit.MMDVAEs.loss), @NamedTuple{}, Matrix{Float32}, Matrix{Float32}}, AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Any}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [21] withgradient(f::Function, args::AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:213
   [22] train!(mmdvae::AutoEncoderToolkit.MMDVAEs.MMDVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, x_in::Matrix{Float32}, x_out::Matrix{Float32}, opt::@NamedTuple{vae::@NamedTuple{encoder::@NamedTuple{encoder::@NamedTuple{layers::Tuple{@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, @NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}}}, μ::@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, logσ::@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}}, decoder::@NamedTuple{decoder::@NamedTuple{layers::Tuple{@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, @NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, @NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}}}}}}; loss_function::typeof(AutoEncoderToolkit.MMDVAEs.loss), loss_kwargs::@NamedTuple{}, verbose::Bool, loss_return::Bool)
      @ AutoEncoderToolkit.MMDVAEs ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/mmdvae.jl:576
   [23] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/mmdvae.jl:291
   [24] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [25] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/mmdvae.jl:318 [inlined]
   [26] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [27] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/mmdvae.jl:329 [inlined]
   [28] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [29] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/mmdvae.jl:330 [inlined]
   [30] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [31] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [32] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [33] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:33 [inlined]
   [34] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [35] top-level scope
      @ none:6
   [36] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [37] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
   [38] _start()
      @ Base ./client.jl:558

All tests passed!


Testing InfoMaxVAEs.jl...

with same input and output: Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/infomaxvae.jl:283
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [3] (::Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [4] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##reparameterize#1", Bool, Type{Float32}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [5] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [6] reparameterize
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:157 [inlined]
    [7] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.reparameterize), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:log,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{log::Bool}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##_#5", Bool, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [9] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [10] #infomaxloss#14
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/infomaxvae.jl:643 [inlined]
   [11] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.InfoMaxVAEs.var"##infomaxloss#14", Float32, Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, Float32, Nothing, typeof(AutoEncoderToolkit.InfoMaxVAEs.infomaxloss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, AutoEncoderToolkit.InfoMaxVAEs.MutualInfoChain, Matrix{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [12] (::Zygote.Pullback{Tuple{var"#64#65"{Matrix{Float32}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:infomaxvae, Zygote.Context{false}, var"#64#65"{Matrix{Float32}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, var"#64#65"{Matrix{Float32}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:mi, Zygote.Context{false}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.InfoMaxVAEs.MutualInfoChain}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.InfoMaxVAEs.infomaxloss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, AutoEncoderToolkit.InfoMaxVAEs.MutualInfoChain, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.InfoMaxVAEs.var"##infomaxloss#14", Float32, Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, Float32, Nothing, typeof(AutoEncoderToolkit.InfoMaxVAEs.infomaxloss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, AutoEncoderToolkit.InfoMaxVAEs.MutualInfoChain, Matrix{Float32}}, Any}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [13] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#64#65"{Matrix{Float32}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:infomaxvae, Zygote.Context{false}, var"#64#65"{Matrix{Float32}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, var"#64#65"{Matrix{Float32}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:mi, Zygote.Context{false}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.InfoMaxVAEs.MutualInfoChain}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.InfoMaxVAEs.infomaxloss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, AutoEncoderToolkit.InfoMaxVAEs.MutualInfoChain, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.InfoMaxVAEs.var"##infomaxloss#14", Float32, Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, Float32, Nothing, typeof(AutoEncoderToolkit.InfoMaxVAEs.infomaxloss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, AutoEncoderToolkit.InfoMaxVAEs.MutualInfoChain, Matrix{Float32}}, Any}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [14] gradient(f::Function, args::AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [15] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/infomaxvae.jl:250
   [16] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [17] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/infomaxvae.jl:284 [inlined]
   [18] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [19] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/infomaxvae.jl:284 [inlined]
   [20] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [21] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [22] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [23] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:36 [inlined]
   [24] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [25] top-level scope
      @ none:6
   [26] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [27] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
   [28] _start()
      @ Base ./client.jl:558
with different input and output: Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/infomaxvae.jl:301
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [3] (::Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [4] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##reparameterize#1", Bool, Type{Float32}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [5] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [6] reparameterize
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:157 [inlined]
    [7] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.reparameterize), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:log,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{log::Bool}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##_#5", Bool, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [9] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [10] #infomaxloss#18
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/infomaxvae.jl:761 [inlined]
   [11] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.InfoMaxVAEs.var"##infomaxloss#18", Float32, Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, Float32, Nothing, typeof(AutoEncoderToolkit.InfoMaxVAEs.infomaxloss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, AutoEncoderToolkit.InfoMaxVAEs.MutualInfoChain, Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [12] (::Zygote.Pullback{Tuple{var"#68#69"{Matrix{Float32}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:infomaxvae, Zygote.Context{false}, var"#68#69"{Matrix{Float32}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, var"#68#69"{Matrix{Float32}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:mi, Zygote.Context{false}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.InfoMaxVAEs.MutualInfoChain}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_out, Zygote.Context{false}, var"#68#69"{Matrix{Float32}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.InfoMaxVAEs.infomaxloss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, AutoEncoderToolkit.InfoMaxVAEs.MutualInfoChain, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.InfoMaxVAEs.var"##infomaxloss#18", Float32, Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, Float32, Nothing, typeof(AutoEncoderToolkit.InfoMaxVAEs.infomaxloss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, AutoEncoderToolkit.InfoMaxVAEs.MutualInfoChain, Matrix{Float32}, Matrix{Float32}}, Any}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [13] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#68#69"{Matrix{Float32}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:infomaxvae, Zygote.Context{false}, var"#68#69"{Matrix{Float32}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, var"#68#69"{Matrix{Float32}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:mi, Zygote.Context{false}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, AutoEncoderToolkit.InfoMaxVAEs.MutualInfoChain}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x_out, Zygote.Context{false}, var"#68#69"{Matrix{Float32}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.InfoMaxVAEs.infomaxloss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, AutoEncoderToolkit.InfoMaxVAEs.MutualInfoChain, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.InfoMaxVAEs.var"##infomaxloss#18", Float32, Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, Float32, Nothing, typeof(AutoEncoderToolkit.InfoMaxVAEs.infomaxloss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, AutoEncoderToolkit.InfoMaxVAEs.MutualInfoChain, Matrix{Float32}, Matrix{Float32}}, Any}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [14] gradient(f::Function, args::AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [15] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/infomaxvae.jl:250
   [16] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [17] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/infomaxvae.jl:302 [inlined]
   [18] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [19] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/infomaxvae.jl:303 [inlined]
   [20] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [21] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [22] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [23] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:36 [inlined]
   [24] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [25] top-level scope
      @ none:6
   [26] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [27] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
   [28] _start()
      @ Base ./client.jl:558
with same input and output: Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/infomaxvae.jl:359
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [3] (::Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [4] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##reparameterize#1", Bool, Type{Float32}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [5] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [6] reparameterize
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:157 [inlined]
    [7] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.reparameterize), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:log,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{log::Bool}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##_#5", Bool, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [9] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [10] #infomaxloss#14
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/infomaxvae.jl:643 [inlined]
   [11] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.InfoMaxVAEs.var"##infomaxloss#14", Float32, Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, Float32, Nothing, typeof(AutoEncoderToolkit.InfoMaxVAEs.infomaxloss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, AutoEncoderToolkit.InfoMaxVAEs.MutualInfoChain, Matrix{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [12] infomaxloss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/infomaxvae.jl:605 [inlined]
   [13] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.InfoMaxVAEs.infomaxloss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, AutoEncoderToolkit.InfoMaxVAEs.MutualInfoChain, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.InfoMaxVAEs.var"##infomaxloss#14", Float32, Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, Float32, Nothing, typeof(AutoEncoderToolkit.InfoMaxVAEs.infomaxloss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, AutoEncoderToolkit.InfoMaxVAEs.MutualInfoChain, Matrix{Float32}}, Any}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [14] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.InfoMaxVAEs.var"#29#30"{typeof(AutoEncoderToolkit.InfoMaxVAEs.infomaxloss), @NamedTuple{}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [15] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{AutoEncoderToolkit.InfoMaxVAEs.var"#29#30"{typeof(AutoEncoderToolkit.InfoMaxVAEs.infomaxloss), @NamedTuple{}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Any}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [16] withgradient(f::Function, args::AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:213
   [17] train!(infomaxvae::AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, x::Matrix{Float32}, opt::@NamedTuple{vae::@NamedTuple{encoder::@NamedTuple{encoder::@NamedTuple{layers::Tuple{@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, @NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}}}, μ::@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, logσ::@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}}, decoder::@NamedTuple{decoder::@NamedTuple{layers::Tuple{@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, @NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, @NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}}}}}, mi::@NamedTuple{data::@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, latent::@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, mlp::@NamedTuple{layers::Tuple{@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, @NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, @NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}}}}}; infomaxloss_function::typeof(AutoEncoderToolkit.InfoMaxVAEs.infomaxloss), infomaxloss_kwargs::@NamedTuple{}, miloss_function::typeof(AutoEncoderToolkit.InfoMaxVAEs.miloss), miloss_kwargs::@NamedTuple{}, loss_return::Bool, verbose::Bool)
      @ AutoEncoderToolkit.InfoMaxVAEs ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/infomaxvae.jl:931
   [18] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/infomaxvae.jl:322
   [19] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [20] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/infomaxvae.jl:354 [inlined]
   [21] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [22] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/infomaxvae.jl:360 [inlined]
   [23] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [24] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/infomaxvae.jl:360 [inlined]
   [25] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [26] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [27] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [28] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:36 [inlined]
   [29] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [30] top-level scope
      @ none:6
   [31] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [32] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
   [33] _start()
      @ Base ./client.jl:558
with different input and output: Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/infomaxvae.jl:366
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [3] (::Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [4] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##reparameterize#1", Bool, Type{Float32}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [5] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [6] reparameterize
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/vae.jl:157 [inlined]
    [7] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.VAEs.reparameterize), JointGaussianLogEncoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:log,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{log::Bool}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{log::Bool}, typeof(AutoEncoderToolkit.VAEs.reparameterize), Matrix{Float32}, Matrix{Float32}}, Any}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.VAEs.var"##_#5", Bool, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
    [9] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{latent::Bool}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, z::Matrix{Float32}, decoder::@NamedTuple{μ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [10] #infomaxloss#18
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/infomaxvae.jl:761 [inlined]
   [11] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.InfoMaxVAEs.var"##infomaxloss#18", Float32, Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, Float32, Nothing, typeof(AutoEncoderToolkit.InfoMaxVAEs.infomaxloss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, AutoEncoderToolkit.InfoMaxVAEs.MutualInfoChain, Matrix{Float32}, Matrix{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [12] infomaxloss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/infomaxvae.jl:722 [inlined]
   [13] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.InfoMaxVAEs.infomaxloss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, AutoEncoderToolkit.InfoMaxVAEs.MutualInfoChain, Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.InfoMaxVAEs.var"##infomaxloss#18", Float32, Float32, typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.encoder_kl), Nothing, Float32, Nothing, typeof(AutoEncoderToolkit.InfoMaxVAEs.infomaxloss), AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}, AutoEncoderToolkit.InfoMaxVAEs.MutualInfoChain, Matrix{Float32}, Matrix{Float32}}, Any}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [14] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.InfoMaxVAEs.var"#36#37"{typeof(AutoEncoderToolkit.InfoMaxVAEs.infomaxloss), @NamedTuple{}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}, Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [15] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{AutoEncoderToolkit.InfoMaxVAEs.var"#36#37"{typeof(AutoEncoderToolkit.InfoMaxVAEs.infomaxloss), @NamedTuple{}, AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Matrix{Float32}, Matrix{Float32}}, AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, Any}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [16] withgradient(f::Function, args::AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:213
   [17] train!(infomaxvae::AutoEncoderToolkit.InfoMaxVAEs.InfoMaxVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, x_in::Matrix{Float32}, x_out::Matrix{Float32}, opt::@NamedTuple{vae::@NamedTuple{encoder::@NamedTuple{encoder::@NamedTuple{layers::Tuple{@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, @NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}}}, μ::@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, logσ::@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}}, decoder::@NamedTuple{decoder::@NamedTuple{layers::Tuple{@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, @NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, @NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}}}}}, mi::@NamedTuple{data::@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, latent::@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, mlp::@NamedTuple{layers::Tuple{@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, @NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, @NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}}}}}; infomaxloss_function::typeof(AutoEncoderToolkit.InfoMaxVAEs.infomaxloss), infomaxloss_kwargs::@NamedTuple{}, miloss_function::typeof(AutoEncoderToolkit.InfoMaxVAEs.miloss), miloss_kwargs::@NamedTuple{}, loss_return::Bool, verbose::Bool)
      @ AutoEncoderToolkit.InfoMaxVAEs ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/infomaxvae.jl:1040
   [18] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/infomaxvae.jl:322
   [19] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [20] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/infomaxvae.jl:354 [inlined]
   [21] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [22] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/infomaxvae.jl:367 [inlined]
   [23] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [24] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/infomaxvae.jl:368 [inlined]
   [25] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [26] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [27] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [28] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:36 [inlined]
   [29] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [30] top-level scope
      @ none:6
   [31] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [32] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
   [33] _start()
      @ Base ./client.jl:558

All tests passed!


Testing HVAEs module:

Generating synthetic data...
WARNING: Method definition f(Any, Any) in module Main at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/vae.jl:68 overwritten at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:42.
Defining layers structure...
Defining encoders...
Defining decoders...
gradient (same input and output): Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:761
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
    [4] unit_vector
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1058 [inlined]
    [5] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [6] #_unit_vectors##0
      @ ./none:-1 [inlined]
    [7] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] #∇map##7
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202 [inlined]
    [9] #Generator##0
      @ ./generator.jl:37 [inlined]
   [10] iterate
      @ ./generator.jl:48 [inlined]
   [11] collect(itr::Base.Generator{Base.Iterators.Zip{Tuple{Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Vector{Vector{Float32}}}}, Base.var"#Generator##0#Generator##1"{Zygote.var"#∇map##7#∇map##8"}})
      @ Base ./array.jl:790
   [12] map
      @ ./abstractarray.jl:3536 [inlined]
   [13] (::Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202
   [14] collect_pullback
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:236 [inlined]
   [15] _unit_vectors
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1127 [inlined]
   [16] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [17] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vectors), Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [20] #finite_difference_gradient#29
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1194 [inlined]
   [21] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##finite_difference_gradient#29", Symbol, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils._finite_difference_gradient), Type{Vector{Float32}}, AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [22] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [23] #∇potential_energy_finite#3
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:293 [inlined]
   [24] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##∇potential_energy_finite#3", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Symbol, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy_finite), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Tuple{Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), typeof(AutoEncoderToolkit.spherical_logprior)}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), JointGaussianLogDecoder{…}}, Any}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), typeof(AutoEncoderToolkit.decoder_loglikelihood)}, Any}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [25] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy_finite), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [26] #∇potential_energy#19
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:565 [inlined]
   [27] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##∇potential_energy#19", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Symbol, @NamedTuple{}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [28] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 3 more times ---
   [32] leapfrog_tempering_step
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:976 [inlined]
   [33] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, ∇U_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering)}, typeof(AutoEncoderToolkit.HVAEs.leapfrog_tempering_step), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Tuple{@NamedTuple{z_init::Vector{Float32}, ρ_init::Vector{Float32}, z_final::Vector{Float32}, ρ_final::Vector{Float32}}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [34] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##_#31", Float32, Int64, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, decoder::@NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, phase_space::@NamedTuple{z_init::Vector{Float32}, ρ_init::Vector{Float32}, z_final::Vector{Float32}, ρ_final::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [36] #hamiltonian_elbo#34
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:1527 [inlined]
   [37] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##hamiltonian_elbo#34", Float32, Int64, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.hamiltonian_elbo), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [38] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, ∇U_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), logp_prefactor::Vector{Float32}, logq_prefactor::Vector{Float32}}, typeof(AutoEncoderToolkit.HVAEs.hamiltonian_elbo), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [40] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:1700 [inlined]
   [41] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#36", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [42] (::Zygote.Pullback{Tuple{var"#72#73", AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#36", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [43] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#72#73", AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#36", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [44] gradient(f::Function, args::AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [45] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:753
   [46] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [47] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:755 [inlined]
   [48] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [49] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:762 [inlined]
  --- the above 2 lines are repeated 1 more time ---
   [52] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [53] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [54] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [55] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:39 [inlined]
   [56] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [57] top-level scope
      @ none:6
   [58] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [59] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
gradient (different input and output): Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:766
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
    [4] unit_vector
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1058 [inlined]
    [5] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [6] #_unit_vectors##0
      @ ./none:-1 [inlined]
    [7] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] #∇map##7
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202 [inlined]
    [9] #Generator##0
      @ ./generator.jl:37 [inlined]
   [10] iterate
      @ ./generator.jl:48 [inlined]
   [11] collect(itr::Base.Generator{Base.Iterators.Zip{Tuple{Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Vector{Vector{Float32}}}}, Base.var"#Generator##0#Generator##1"{Zygote.var"#∇map##7#∇map##8"}})
      @ Base ./array.jl:790
   [12] map
      @ ./abstractarray.jl:3536 [inlined]
   [13] (::Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202
   [14] collect_pullback
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:236 [inlined]
   [15] _unit_vectors
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1127 [inlined]
   [16] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [17] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vectors), Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [20] #finite_difference_gradient#29
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1194 [inlined]
   [21] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##finite_difference_gradient#29", Symbol, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils._finite_difference_gradient), Type{Vector{Float32}}, AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [22] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [23] #∇potential_energy_finite#3
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:293 [inlined]
   [24] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##∇potential_energy_finite#3", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Symbol, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy_finite), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Tuple{Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), typeof(AutoEncoderToolkit.spherical_logprior)}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), JointGaussianLogDecoder{…}}, Any}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), typeof(AutoEncoderToolkit.decoder_loglikelihood)}, Any}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [25] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy_finite), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [26] #∇potential_energy#19
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:565 [inlined]
   [27] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##∇potential_energy#19", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Symbol, @NamedTuple{}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [28] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 3 more times ---
   [32] leapfrog_tempering_step
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:976 [inlined]
   [33] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, ∇U_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering)}, typeof(AutoEncoderToolkit.HVAEs.leapfrog_tempering_step), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Tuple{@NamedTuple{z_init::Vector{Float32}, ρ_init::Vector{Float32}, z_final::Vector{Float32}, ρ_final::Vector{Float32}}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [34] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##_#31", Float32, Int64, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, decoder::@NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, phase_space::@NamedTuple{z_init::Vector{Float32}, ρ_init::Vector{Float32}, z_final::Vector{Float32}, ρ_final::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [36] #hamiltonian_elbo#35
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:1640 [inlined]
   [37] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##hamiltonian_elbo#35", Float32, Int64, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.hamiltonian_elbo), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [38] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, ∇U_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), logp_prefactor::Vector{Float32}, logq_prefactor::Vector{Float32}}, typeof(AutoEncoderToolkit.HVAEs.hamiltonian_elbo), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [40] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:1804 [inlined]
   [41] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#37", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [42] (::Zygote.Pullback{Tuple{var"#74#75", AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#37", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [43] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#74#75", AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#37", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [44] gradient(f::Function, args::AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [45] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:753
   [46] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [47] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:755 [inlined]
   [48] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [49] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:767 [inlined]
  --- the above 2 lines are repeated 1 more time ---
   [52] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [53] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [54] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [55] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:39 [inlined]
   [56] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [57] top-level scope
      @ none:6
   [58] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [59] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
gradient (same input and output): Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:761
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
    [4] unit_vector
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1058 [inlined]
    [5] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [6] #_unit_vectors##0
      @ ./none:-1 [inlined]
    [7] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] #∇map##7
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202 [inlined]
    [9] #Generator##0
      @ ./generator.jl:37 [inlined]
   [10] iterate
      @ ./generator.jl:48 [inlined]
   [11] collect(itr::Base.Generator{Base.Iterators.Zip{Tuple{Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Vector{Vector{Float32}}}}, Base.var"#Generator##0#Generator##1"{Zygote.var"#∇map##7#∇map##8"}})
      @ Base ./array.jl:790
   [12] map
      @ ./abstractarray.jl:3536 [inlined]
   [13] (::Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202
   [14] collect_pullback
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:236 [inlined]
   [15] _unit_vectors
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1127 [inlined]
   [16] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [17] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vectors), Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [20] #finite_difference_gradient#29
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1194 [inlined]
   [21] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##finite_difference_gradient#29", Symbol, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils._finite_difference_gradient), Type{Vector{Float32}}, AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [22] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [23] #∇potential_energy_finite#3
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:293 [inlined]
   [24] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##∇potential_energy_finite#3", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Symbol, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy_finite), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Tuple{Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), typeof(AutoEncoderToolkit.spherical_logprior)}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), SplitGaussianLogDecoder{…}}, Any}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), typeof(AutoEncoderToolkit.decoder_loglikelihood)}, Any}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [25] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy_finite), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [26] #∇potential_energy#19
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:565 [inlined]
   [27] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##∇potential_energy#19", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Symbol, @NamedTuple{}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [28] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 3 more times ---
   [32] leapfrog_tempering_step
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:976 [inlined]
   [33] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, ∇U_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering)}, typeof(AutoEncoderToolkit.HVAEs.leapfrog_tempering_step), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Tuple{@NamedTuple{z_init::Vector{Float32}, ρ_init::Vector{Float32}, z_final::Vector{Float32}, ρ_final::Vector{Float32}}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [34] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##_#31", Float32, Int64, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, decoder::@NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, phase_space::@NamedTuple{z_init::Vector{Float32}, ρ_init::Vector{Float32}, z_final::Vector{Float32}, ρ_final::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [36] #hamiltonian_elbo#34
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:1527 [inlined]
   [37] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##hamiltonian_elbo#34", Float32, Int64, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.hamiltonian_elbo), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [38] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, ∇U_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), logp_prefactor::Vector{Float32}, logq_prefactor::Vector{Float32}}, typeof(AutoEncoderToolkit.HVAEs.hamiltonian_elbo), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [40] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:1700 [inlined]
   [41] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#36", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [42] (::Zygote.Pullback{Tuple{var"#72#73", AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#36", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [43] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#72#73", AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#36", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [44] gradient(f::Function, args::AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [45] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:753
   [46] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [47] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:755 [inlined]
   [48] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [49] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:762 [inlined]
  --- the above 2 lines are repeated 1 more time ---
   [52] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [53] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [54] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [55] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:39 [inlined]
   [56] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [57] top-level scope
      @ none:6
   [58] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [59] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
gradient (different input and output): Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:766
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
    [4] unit_vector
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1058 [inlined]
    [5] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [6] #_unit_vectors##0
      @ ./none:-1 [inlined]
    [7] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] #∇map##7
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202 [inlined]
    [9] #Generator##0
      @ ./generator.jl:37 [inlined]
   [10] iterate
      @ ./generator.jl:48 [inlined]
   [11] collect(itr::Base.Generator{Base.Iterators.Zip{Tuple{Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Vector{Vector{Float32}}}}, Base.var"#Generator##0#Generator##1"{Zygote.var"#∇map##7#∇map##8"}})
      @ Base ./array.jl:790
   [12] map
      @ ./abstractarray.jl:3536 [inlined]
   [13] (::Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202
   [14] collect_pullback
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:236 [inlined]
   [15] _unit_vectors
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1127 [inlined]
   [16] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [17] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vectors), Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [20] #finite_difference_gradient#29
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1194 [inlined]
   [21] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##finite_difference_gradient#29", Symbol, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils._finite_difference_gradient), Type{Vector{Float32}}, AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [22] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [23] #∇potential_energy_finite#3
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:293 [inlined]
   [24] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##∇potential_energy_finite#3", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Symbol, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy_finite), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Tuple{Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), typeof(AutoEncoderToolkit.spherical_logprior)}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), SplitGaussianLogDecoder{…}}, Any}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), typeof(AutoEncoderToolkit.decoder_loglikelihood)}, Any}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [25] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy_finite), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [26] #∇potential_energy#19
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:565 [inlined]
   [27] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##∇potential_energy#19", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Symbol, @NamedTuple{}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [28] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 3 more times ---
   [32] leapfrog_tempering_step
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:976 [inlined]
   [33] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, ∇U_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering)}, typeof(AutoEncoderToolkit.HVAEs.leapfrog_tempering_step), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Tuple{@NamedTuple{z_init::Vector{Float32}, ρ_init::Vector{Float32}, z_final::Vector{Float32}, ρ_final::Vector{Float32}}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [34] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##_#31", Float32, Int64, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, decoder::@NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, phase_space::@NamedTuple{z_init::Vector{Float32}, ρ_init::Vector{Float32}, z_final::Vector{Float32}, ρ_final::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [36] #hamiltonian_elbo#35
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:1640 [inlined]
   [37] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##hamiltonian_elbo#35", Float32, Int64, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.hamiltonian_elbo), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [38] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, ∇U_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), logp_prefactor::Vector{Float32}, logq_prefactor::Vector{Float32}}, typeof(AutoEncoderToolkit.HVAEs.hamiltonian_elbo), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [40] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:1804 [inlined]
   [41] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#37", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [42] (::Zygote.Pullback{Tuple{var"#74#75", AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#37", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [43] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#74#75", AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#37", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [44] gradient(f::Function, args::AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [45] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:753
   [46] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [47] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:755 [inlined]
   [48] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [49] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:767 [inlined]
  --- the above 2 lines are repeated 1 more time ---
   [52] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [53] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [54] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [55] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:39 [inlined]
   [56] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [57] top-level scope
      @ none:6
   [58] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [59] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
gradient (same input and output): Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:761
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
    [4] unit_vector
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1058 [inlined]
    [5] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [6] #_unit_vectors##0
      @ ./none:-1 [inlined]
    [7] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] #∇map##7
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202 [inlined]
    [9] #Generator##0
      @ ./generator.jl:37 [inlined]
   [10] iterate
      @ ./generator.jl:48 [inlined]
   [11] collect(itr::Base.Generator{Base.Iterators.Zip{Tuple{Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Vector{Vector{Float32}}}}, Base.var"#Generator##0#Generator##1"{Zygote.var"#∇map##7#∇map##8"}})
      @ Base ./array.jl:790
   [12] map
      @ ./abstractarray.jl:3536 [inlined]
   [13] (::Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202
   [14] collect_pullback
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:236 [inlined]
   [15] _unit_vectors
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1127 [inlined]
   [16] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [17] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vectors), Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [20] #finite_difference_gradient#29
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1194 [inlined]
   [21] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##finite_difference_gradient#29", Symbol, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils._finite_difference_gradient), Type{Vector{Float32}}, AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [22] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [23] #∇potential_energy_finite#3
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:293 [inlined]
   [24] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##∇potential_energy_finite#3", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Symbol, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy_finite), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Tuple{Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), typeof(AutoEncoderToolkit.spherical_logprior)}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), JointGaussianDecoder{…}}, Any}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), typeof(AutoEncoderToolkit.decoder_loglikelihood)}, Any}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [25] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy_finite), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [26] #∇potential_energy#19
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:565 [inlined]
   [27] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##∇potential_energy#19", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Symbol, @NamedTuple{}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [28] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 3 more times ---
   [32] leapfrog_tempering_step
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:976 [inlined]
   [33] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, ∇U_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering)}, typeof(AutoEncoderToolkit.HVAEs.leapfrog_tempering_step), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any})(Δ::Tuple{@NamedTuple{z_init::Vector{Float32}, ρ_init::Vector{Float32}, z_final::Vector{Float32}, ρ_final::Vector{Float32}}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [34] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##_#31", Float32, Int64, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, decoder::@NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}, phase_space::@NamedTuple{z_init::Vector{Float32}, ρ_init::Vector{Float32}, z_final::Vector{Float32}, ρ_final::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [36] #hamiltonian_elbo#34
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:1527 [inlined]
   [37] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##hamiltonian_elbo#34", Float32, Int64, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.hamiltonian_elbo), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [38] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, ∇U_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), logp_prefactor::Vector{Float32}, logq_prefactor::Vector{Float32}}, typeof(AutoEncoderToolkit.HVAEs.hamiltonian_elbo), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [40] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:1700 [inlined]
   [41] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#36", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [42] (::Zygote.Pullback{Tuple{var"#72#73", AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#36", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [43] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#72#73", AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#36", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [44] gradient(f::Function, args::AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [45] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:753
   [46] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [47] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:755 [inlined]
   [48] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [49] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:762 [inlined]
  --- the above 2 lines are repeated 1 more time ---
   [52] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [53] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [54] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [55] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:39 [inlined]
   [56] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [57] top-level scope
      @ none:6
   [58] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [59] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
gradient (different input and output): Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:766
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
    [4] unit_vector
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1058 [inlined]
    [5] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [6] #_unit_vectors##0
      @ ./none:-1 [inlined]
    [7] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] #∇map##7
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202 [inlined]
    [9] #Generator##0
      @ ./generator.jl:37 [inlined]
   [10] iterate
      @ ./generator.jl:48 [inlined]
   [11] collect(itr::Base.Generator{Base.Iterators.Zip{Tuple{Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Vector{Vector{Float32}}}}, Base.var"#Generator##0#Generator##1"{Zygote.var"#∇map##7#∇map##8"}})
      @ Base ./array.jl:790
   [12] map
      @ ./abstractarray.jl:3536 [inlined]
   [13] (::Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202
   [14] collect_pullback
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:236 [inlined]
   [15] _unit_vectors
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1127 [inlined]
   [16] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [17] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vectors), Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [20] #finite_difference_gradient#29
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1194 [inlined]
   [21] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##finite_difference_gradient#29", Symbol, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils._finite_difference_gradient), Type{Vector{Float32}}, AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [22] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [23] #∇potential_energy_finite#3
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:293 [inlined]
   [24] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##∇potential_energy_finite#3", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Symbol, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy_finite), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Tuple{Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), typeof(AutoEncoderToolkit.spherical_logprior)}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), JointGaussianDecoder{…}}, Any}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), typeof(AutoEncoderToolkit.decoder_loglikelihood)}, Any}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [25] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy_finite), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [26] #∇potential_energy#19
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:565 [inlined]
   [27] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##∇potential_energy#19", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Symbol, @NamedTuple{}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [28] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 3 more times ---
   [32] leapfrog_tempering_step
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:976 [inlined]
   [33] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, ∇U_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering)}, typeof(AutoEncoderToolkit.HVAEs.leapfrog_tempering_step), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any})(Δ::Tuple{@NamedTuple{z_init::Vector{Float32}, ρ_init::Vector{Float32}, z_final::Vector{Float32}, ρ_final::Vector{Float32}}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [34] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##_#31", Float32, Int64, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, decoder::@NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}, phase_space::@NamedTuple{z_init::Vector{Float32}, ρ_init::Vector{Float32}, z_final::Vector{Float32}, ρ_final::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [36] #hamiltonian_elbo#35
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:1640 [inlined]
   [37] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##hamiltonian_elbo#35", Float32, Int64, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.hamiltonian_elbo), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [38] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, ∇U_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), logp_prefactor::Vector{Float32}, logq_prefactor::Vector{Float32}}, typeof(AutoEncoderToolkit.HVAEs.hamiltonian_elbo), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [40] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:1804 [inlined]
   [41] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#37", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [42] (::Zygote.Pullback{Tuple{var"#74#75", AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#37", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [43] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#74#75", AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#37", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [44] gradient(f::Function, args::AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [45] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:753
   [46] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [47] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:755 [inlined]
   [48] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [49] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:767 [inlined]
  --- the above 2 lines are repeated 1 more time ---
   [52] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [53] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [54] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [55] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:39 [inlined]
   [56] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [57] top-level scope
      @ none:6
   [58] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [59] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
gradient (same input and output): Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:761
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
    [4] unit_vector
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1058 [inlined]
    [5] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [6] #_unit_vectors##0
      @ ./none:-1 [inlined]
    [7] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] #∇map##7
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202 [inlined]
    [9] #Generator##0
      @ ./generator.jl:37 [inlined]
   [10] iterate
      @ ./generator.jl:48 [inlined]
   [11] collect(itr::Base.Generator{Base.Iterators.Zip{Tuple{Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Vector{Vector{Float32}}}}, Base.var"#Generator##0#Generator##1"{Zygote.var"#∇map##7#∇map##8"}})
      @ Base ./array.jl:790
   [12] map
      @ ./abstractarray.jl:3536 [inlined]
   [13] (::Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202
   [14] collect_pullback
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:236 [inlined]
   [15] _unit_vectors
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1127 [inlined]
   [16] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [17] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vectors), Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [20] #finite_difference_gradient#29
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1194 [inlined]
   [21] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##finite_difference_gradient#29", Symbol, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils._finite_difference_gradient), Type{Vector{Float32}}, AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [22] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [23] #∇potential_energy_finite#3
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:293 [inlined]
   [24] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##∇potential_energy_finite#3", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Symbol, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy_finite), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Tuple{Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), typeof(AutoEncoderToolkit.spherical_logprior)}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), SplitGaussianDecoder{…}}, Any}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), typeof(AutoEncoderToolkit.decoder_loglikelihood)}, Any}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [25] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy_finite), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [26] #∇potential_energy#19
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:565 [inlined]
   [27] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##∇potential_energy#19", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Symbol, @NamedTuple{}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [28] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 3 more times ---
   [32] leapfrog_tempering_step
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:976 [inlined]
   [33] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, ∇U_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering)}, typeof(AutoEncoderToolkit.HVAEs.leapfrog_tempering_step), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any})(Δ::Tuple{@NamedTuple{z_init::Vector{Float32}, ρ_init::Vector{Float32}, z_final::Vector{Float32}, ρ_final::Vector{Float32}}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [34] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##_#31", Float32, Int64, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, decoder::@NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}, phase_space::@NamedTuple{z_init::Vector{Float32}, ρ_init::Vector{Float32}, z_final::Vector{Float32}, ρ_final::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [36] #hamiltonian_elbo#34
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:1527 [inlined]
   [37] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##hamiltonian_elbo#34", Float32, Int64, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.hamiltonian_elbo), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [38] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, ∇U_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), logp_prefactor::Vector{Float32}, logq_prefactor::Vector{Float32}}, typeof(AutoEncoderToolkit.HVAEs.hamiltonian_elbo), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [40] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:1700 [inlined]
   [41] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#36", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [42] (::Zygote.Pullback{Tuple{var"#72#73", AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#36", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [43] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#72#73", AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#36", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [44] gradient(f::Function, args::AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [45] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:753
   [46] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [47] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:755 [inlined]
   [48] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [49] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:762 [inlined]
  --- the above 2 lines are repeated 1 more time ---
   [52] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [53] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [54] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [55] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:39 [inlined]
   [56] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [57] top-level scope
      @ none:6
   [58] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [59] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
gradient (different input and output): Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:766
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
    [4] unit_vector
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1058 [inlined]
    [5] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [6] #_unit_vectors##0
      @ ./none:-1 [inlined]
    [7] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] #∇map##7
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202 [inlined]
    [9] #Generator##0
      @ ./generator.jl:37 [inlined]
   [10] iterate
      @ ./generator.jl:48 [inlined]
   [11] collect(itr::Base.Generator{Base.Iterators.Zip{Tuple{Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Vector{Vector{Float32}}}}, Base.var"#Generator##0#Generator##1"{Zygote.var"#∇map##7#∇map##8"}})
      @ Base ./array.jl:790
   [12] map
      @ ./abstractarray.jl:3536 [inlined]
   [13] (::Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202
   [14] collect_pullback
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:236 [inlined]
   [15] _unit_vectors
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1127 [inlined]
   [16] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [17] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vectors), Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [20] #finite_difference_gradient#29
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1194 [inlined]
   [21] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##finite_difference_gradient#29", Symbol, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils._finite_difference_gradient), Type{Vector{Float32}}, AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [22] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [23] #∇potential_energy_finite#3
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:293 [inlined]
   [24] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##∇potential_energy_finite#3", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Symbol, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy_finite), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Tuple{Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), typeof(AutoEncoderToolkit.spherical_logprior)}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), SplitGaussianDecoder{…}}, Any}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), typeof(AutoEncoderToolkit.decoder_loglikelihood)}, Any}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [25] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy_finite), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [26] #∇potential_energy#19
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:565 [inlined]
   [27] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##∇potential_energy#19", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Symbol, @NamedTuple{}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [28] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 3 more times ---
   [32] leapfrog_tempering_step
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:976 [inlined]
   [33] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, ∇U_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering)}, typeof(AutoEncoderToolkit.HVAEs.leapfrog_tempering_step), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any})(Δ::Tuple{@NamedTuple{z_init::Vector{Float32}, ρ_init::Vector{Float32}, z_final::Vector{Float32}, ρ_final::Vector{Float32}}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [34] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##_#31", Float32, Int64, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, decoder::@NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}, phase_space::@NamedTuple{z_init::Vector{Float32}, ρ_init::Vector{Float32}, z_final::Vector{Float32}, ρ_final::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [36] #hamiltonian_elbo#35
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:1640 [inlined]
   [37] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##hamiltonian_elbo#35", Float32, Int64, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.hamiltonian_elbo), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [38] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, ∇U_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), logp_prefactor::Vector{Float32}, logq_prefactor::Vector{Float32}}, typeof(AutoEncoderToolkit.HVAEs.hamiltonian_elbo), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [40] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:1804 [inlined]
   [41] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#37", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [42] (::Zygote.Pullback{Tuple{var"#74#75", AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#37", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [43] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#74#75", AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#37", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [44] gradient(f::Function, args::AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianDecoder{…}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [45] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:753
   [46] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [47] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:755 [inlined]
   [48] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [49] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:767 [inlined]
  --- the above 2 lines are repeated 1 more time ---
   [52] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [53] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [54] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [55] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:39 [inlined]
   [56] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [57] top-level scope
      @ none:6
   [58] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [59] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
gradient (same input and output): Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:761
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
    [4] unit_vector
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1058 [inlined]
    [5] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [6] #_unit_vectors##0
      @ ./none:-1 [inlined]
    [7] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] #∇map##7
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202 [inlined]
    [9] #Generator##0
      @ ./generator.jl:37 [inlined]
   [10] iterate
      @ ./generator.jl:48 [inlined]
   [11] collect(itr::Base.Generator{Base.Iterators.Zip{Tuple{Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Vector{Vector{Float32}}}}, Base.var"#Generator##0#Generator##1"{Zygote.var"#∇map##7#∇map##8"}})
      @ Base ./array.jl:790
   [12] map
      @ ./abstractarray.jl:3536 [inlined]
   [13] (::Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202
   [14] collect_pullback
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:236 [inlined]
   [15] _unit_vectors
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1127 [inlined]
   [16] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [17] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vectors), Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [20] #finite_difference_gradient#29
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1194 [inlined]
   [21] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##finite_difference_gradient#29", Symbol, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils._finite_difference_gradient), Type{Vector{Float32}}, AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}}}, Vector{Float32}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [22] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}}}, Vector{Float32}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [23] #∇potential_energy_finite#3
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:293 [inlined]
   [24] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##∇potential_energy_finite#3", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Symbol, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy_finite), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}}}, Tuple{Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), typeof(AutoEncoderToolkit.spherical_logprior)}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), SimpleGaussianDecoder{…}}, Any}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), typeof(AutoEncoderToolkit.decoder_loglikelihood)}, Any}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}}}, Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), @NamedTuple{μ::Vector{Float32}}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [25] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy_finite), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [26] #∇potential_energy#19
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:565 [inlined]
   [27] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##∇potential_energy#19", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Symbol, @NamedTuple{}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [28] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 3 more times ---
   [32] leapfrog_tempering_step
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:976 [inlined]
   [33] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, ∇U_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering)}, typeof(AutoEncoderToolkit.HVAEs.leapfrog_tempering_step), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}}}, Any})(Δ::Tuple{@NamedTuple{z_init::Vector{Float32}, ρ_init::Vector{Float32}, z_final::Vector{Float32}, ρ_final::Vector{Float32}}, @NamedTuple{μ::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [34] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##_#31", Float32, Int64, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, decoder::@NamedTuple{μ::Vector{Float32}}, phase_space::@NamedTuple{z_init::Vector{Float32}, ρ_init::Vector{Float32}, z_final::Vector{Float32}, ρ_final::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [36] #hamiltonian_elbo#34
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:1527 [inlined]
   [37] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##hamiltonian_elbo#34", Float32, Int64, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.hamiltonian_elbo), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [38] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, ∇U_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), logp_prefactor::Vector{Float32}, logq_prefactor::Vector{Float32}}, typeof(AutoEncoderToolkit.HVAEs.hamiltonian_elbo), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [40] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:1700 [inlined]
   [41] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#36", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [42] (::Zygote.Pullback{Tuple{var"#72#73", AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#36", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [43] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#72#73", AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#36", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [44] gradient(f::Function, args::AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [45] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:753
   [46] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [47] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:755 [inlined]
   [48] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [49] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:762 [inlined]
  --- the above 2 lines are repeated 1 more time ---
   [52] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [53] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [54] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [55] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:39 [inlined]
   [56] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [57] top-level scope
      @ none:6
   [58] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [59] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
gradient (different input and output): Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:766
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
    [4] unit_vector
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1058 [inlined]
    [5] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [6] #_unit_vectors##0
      @ ./none:-1 [inlined]
    [7] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] #∇map##7
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202 [inlined]
    [9] #Generator##0
      @ ./generator.jl:37 [inlined]
   [10] iterate
      @ ./generator.jl:48 [inlined]
   [11] collect(itr::Base.Generator{Base.Iterators.Zip{Tuple{Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Vector{Vector{Float32}}}}, Base.var"#Generator##0#Generator##1"{Zygote.var"#∇map##7#∇map##8"}})
      @ Base ./array.jl:790
   [12] map
      @ ./abstractarray.jl:3536 [inlined]
   [13] (::Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202
   [14] collect_pullback
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:236 [inlined]
   [15] _unit_vectors
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1127 [inlined]
   [16] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [17] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vectors), Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [20] #finite_difference_gradient#29
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1194 [inlined]
   [21] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##finite_difference_gradient#29", Symbol, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils._finite_difference_gradient), Type{Vector{Float32}}, AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}}}, Vector{Float32}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [22] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}}}, Vector{Float32}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [23] #∇potential_energy_finite#3
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:293 [inlined]
   [24] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##∇potential_energy_finite#3", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Symbol, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy_finite), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}}}, Tuple{Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), typeof(AutoEncoderToolkit.spherical_logprior)}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), SimpleGaussianDecoder{…}}, Any}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), typeof(AutoEncoderToolkit.decoder_loglikelihood)}, Any}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}}}, Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), @NamedTuple{μ::Vector{Float32}}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [25] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy_finite), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [26] #∇potential_energy#19
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:565 [inlined]
   [27] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##∇potential_energy#19", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Symbol, @NamedTuple{}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [28] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 3 more times ---
   [32] leapfrog_tempering_step
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:976 [inlined]
   [33] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, ∇U_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering)}, typeof(AutoEncoderToolkit.HVAEs.leapfrog_tempering_step), SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, Vector{Float32}, SimpleGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}}}, Any})(Δ::Tuple{@NamedTuple{z_init::Vector{Float32}, ρ_init::Vector{Float32}, z_final::Vector{Float32}, ρ_final::Vector{Float32}}, @NamedTuple{μ::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [34] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##_#31", Float32, Int64, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, decoder::@NamedTuple{μ::Vector{Float32}}, phase_space::@NamedTuple{z_init::Vector{Float32}, ρ_init::Vector{Float32}, z_final::Vector{Float32}, ρ_final::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [36] #hamiltonian_elbo#35
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:1640 [inlined]
   [37] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##hamiltonian_elbo#35", Float32, Int64, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.hamiltonian_elbo), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [38] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, ∇U_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), logp_prefactor::Vector{Float32}, logq_prefactor::Vector{Float32}}, typeof(AutoEncoderToolkit.HVAEs.hamiltonian_elbo), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [40] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:1804 [inlined]
   [41] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#37", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [42] (::Zygote.Pullback{Tuple{var"#74#75", AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#37", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [43] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#74#75", AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#37", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, SubArray{Float32, 1, Matrix{Float32}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [44] gradient(f::Function, args::AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SimpleGaussianDecoder{…}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [45] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:753
   [46] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [47] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:755 [inlined]
   [48] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [49] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:767 [inlined]
  --- the above 2 lines are repeated 1 more time ---
   [52] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [53] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [54] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [55] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:39 [inlined]
   [56] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [57] top-level scope
      @ none:6
   [58] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [59] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
CPU | without regularization: Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:785
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
    [4] unit_vector
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1084 [inlined]
    [5] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Matrix{Float32}, Int64}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, 1, Tuple{Base.OneTo{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(eltype), Matrix{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, Base.OneTo{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{Base.OneTo{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}}, Nothing, false}}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [6] #_unit_vectors##2
      @ ./none:-1 [inlined]
    [7] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#fill_pullback#rrule##121"{Tuple{ChainRulesCore.NoTangent}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Matrix{Float32}, Int64}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, 1, Tuple{Base.OneTo{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(eltype), Matrix{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, Base.OneTo{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{Base.OneTo{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}}, Nothing, false}}}}}}, Zygote.ZBack{ChainRules.var"#reduce_hcat_pullback_1#rrule##94"{Base.OneTo{Int64}}}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] #∇map##7
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202 [inlined]
    [9] #Generator##0
      @ ./generator.jl:37 [inlined]
   [10] iterate
      @ ./generator.jl:48 [inlined]
   [11] collect(itr::Base.Generator{Base.Iterators.Zip{Tuple{Vector{Tuple{Matrix{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#fill_pullback#rrule##121"{Tuple{ChainRulesCore.NoTangent}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Matrix{Float32}, Int64}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, 1, Tuple{Base.OneTo{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(eltype), Matrix{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, Base.OneTo{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{Base.OneTo{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}}, Nothing, false}}}}}}, Zygote.ZBack{ChainRules.var"#reduce_hcat_pullback_1#rrule##94"{Base.OneTo{Int64}}}}}}}, Vector{Matrix{Float32}}}}, Base.var"#Generator##0#Generator##1"{Zygote.var"#∇map##7#∇map##8"}})
      @ Base ./array.jl:790
   [12] map
      @ ./abstractarray.jl:3536 [inlined]
   [13] (::Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Matrix{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#fill_pullback#rrule##121"{Tuple{ChainRulesCore.NoTangent}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Matrix{Float32}, Int64}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, 1, Tuple{Base.OneTo{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(eltype), Matrix{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, Base.OneTo{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{Base.OneTo{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}}, Nothing, false}}}}}}, Zygote.ZBack{ChainRules.var"#reduce_hcat_pullback_1#rrule##94"{Base.OneTo{Int64}}}}}}}})(Δ::Vector{Matrix{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202
   [14] collect_pullback
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:236 [inlined]
   [15] _unit_vectors
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1138 [inlined]
   [16] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Matrix{Float32}}, Matrix{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Matrix{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Matrix{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#fill_pullback#rrule##121"{Tuple{ChainRulesCore.NoTangent}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Matrix{Float32}, Int64}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, 1, Tuple{Base.OneTo{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(eltype), Matrix{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, Base.OneTo{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{Base.OneTo{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}}, Nothing, false}}}}}}, Zygote.ZBack{ChainRules.var"#reduce_hcat_pullback_1#rrule##94"{Base.OneTo{Int64}}}}}}}}, Nothing}}})(Δ::Vector{Matrix{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [17] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vectors), Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Matrix{Float32}}, Matrix{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Matrix{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Matrix{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##2#_unit_vectors##3"{Matrix{Float32}}, Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#fill_pullback#rrule##121"{Tuple{ChainRulesCore.NoTangent}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Matrix{Float32}, Int64}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, 1, Tuple{Base.OneTo{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(eltype), Matrix{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}, Base.OneTo{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{Base.OneTo{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##2#unit_vector##3"{Int64, Type{Float32}}}, Nothing, false}}}}}}, Zygote.ZBack{ChainRules.var"#reduce_hcat_pullback_1#rrule##94"{Base.OneTo{Int64}}}}}}}}, Nothing}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Matrix{Float32}}, Any}}})(Δ::Vector{Matrix{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [20] #finite_difference_gradient#29
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1194 [inlined]
   [21] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##finite_difference_gradient#29", Symbol, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Matrix{Float32}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Matrix{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils._finite_difference_gradient), Type{Matrix{Float32}}, AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Matrix{Float32}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Matrix{Float32}}, Any}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [22] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Matrix{Float32}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Matrix{Float32}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [23] #∇potential_energy_finite#3
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:293 [inlined]
   [24] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##∇potential_energy_finite#3", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Symbol, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy_finite), Matrix{Float32}, Matrix{Float32}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Tuple{Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), typeof(AutoEncoderToolkit.spherical_logprior)}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), JointGaussianLogDecoder{…}}, Any}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Matrix{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Matrix{Float32}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), typeof(AutoEncoderToolkit.decoder_loglikelihood)}, Any}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.HVAEs.var"#5#6"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Matrix{Float32}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Matrix{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Any}}})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [25] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy_finite), Matrix{Float32}, Matrix{Float32}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [26] #∇potential_energy#19
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:565 [inlined]
   [27] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##∇potential_energy#19", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), Symbol, @NamedTuple{}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy), Matrix{Float32}, Matrix{Float32}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [28] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.∇potential_energy), Matrix{Float32}, Matrix{Float32}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Any})(Δ::Matrix{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 3 more times ---
   [32] leapfrog_tempering_step
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:976 [inlined]
   [33] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, ∇U_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering)}, typeof(AutoEncoderToolkit.HVAEs.leapfrog_tempering_step), Matrix{Float32}, Matrix{Float32}, JointGaussianLogDecoder{…}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}}, Any})(Δ::Tuple{@NamedTuple{z_init::Matrix{Float32}, ρ_init::Matrix{Float32}, z_final::Matrix{Float32}, ρ_final::Matrix{Float32}}, @NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [34] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##_#31", Float32, Int64, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{encoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, decoder::@NamedTuple{μ::Matrix{Float32}, logσ::Matrix{Float32}}, phase_space::@NamedTuple{z_init::Matrix{Float32}, ρ_init::Matrix{Float32}, z_final::Matrix{Float32}, ρ_final::Matrix{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [36] #hamiltonian_elbo#34
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:1527 [inlined]
   [37] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##hamiltonian_elbo#34", Float32, Int64, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.hamiltonian_elbo), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, Matrix{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [38] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, ∇U_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), logp_prefactor::Vector{Float32}, logq_prefactor::Vector{Float32}}, typeof(AutoEncoderToolkit.HVAEs.hamiltonian_elbo), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, Matrix{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [40] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:1700 [inlined]
   [41] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"##loss#36", Int64, Float32, Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.HVAEs.loss), AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, Matrix{Float32}}, Any}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :latent_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), latent_logprior::typeof(AutoEncoderToolkit.spherical_logprior)}, Nothing, true}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [42] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"#40#41"{typeof(AutoEncoderToolkit.HVAEs.loss), @NamedTuple{}, Matrix{Float32}}, AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [43] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{AutoEncoderToolkit.HVAEs.var"#40#41"{typeof(AutoEncoderToolkit.HVAEs.loss), @NamedTuple{}, Matrix{Float32}}, AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}}, Any}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [44] withgradient(f::Function, args::AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:213
   [45] train!(hvae::AutoEncoderToolkit.HVAEs.HVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, x::Matrix{Float32}, opt::@NamedTuple{vae::@NamedTuple{encoder::@NamedTuple{encoder::@NamedTuple{layers::Tuple{@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, @NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}}}, μ::@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, logσ::@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}}, decoder::@NamedTuple{decoder::@NamedTuple{layers::Tuple{@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, @NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}}}, μ::@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}, logσ::@NamedTuple{weight::Optimisers.Leaf{Optimisers.Adam, Tuple{Matrix{Float32}, Matrix{Float32}, Tuple{Float32, Float32}}}, bias::Optimisers.Leaf{Optimisers.Adam, Tuple{Vector{Float32}, Vector{Float32}, Tuple{Float32, Float32}}}, σ::Tuple{}}}}}; loss_function::typeof(AutoEncoderToolkit.HVAEs.loss), loss_kwargs::@NamedTuple{}, verbose::Bool, loss_return::Bool)
      @ AutoEncoderToolkit.HVAEs ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/hvae.jl:1904
   [46] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:785
   [47] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [48] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:787 [inlined]
   [49] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [50] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:800 [inlined]
   [51] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [52] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [53] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [54] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:39 [inlined]
   [55] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [56] top-level scope
      @ none:6
   [57] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [58] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295

All tests passed!


Testing RHVAEs module:

Generating synthetic data...
WARNING: Method definition f(Any, Any) in module Main at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/hvae.jl:42 overwritten at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:48.
Defining layers structure...
Defining encoders...
Defining decoders...
gradient (same input and output): Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1557
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
    [4] unit_vector
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1058 [inlined]
    [5] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [6] #_unit_vectors##0
      @ ./none:-1 [inlined]
    [7] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] #∇map##7
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202 [inlined]
    [9] #Generator##0
      @ ./generator.jl:37 [inlined]
   [10] iterate
      @ ./generator.jl:48 [inlined]
   [11] collect(itr::Base.Generator{Base.Iterators.Zip{Tuple{Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Vector{Vector{Float32}}}}, Base.var"#Generator##0#Generator##1"{Zygote.var"#∇map##7#∇map##8"}})
      @ Base ./array.jl:790
   [12] map
      @ ./abstractarray.jl:3536 [inlined]
   [13] (::Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202
   [14] collect_pullback
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:236 [inlined]
   [15] _unit_vectors
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1127 [inlined]
   [16] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [17] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vectors), Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [20] #finite_difference_gradient#29
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1194 [inlined]
   [21] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##finite_difference_gradient#29", Symbol, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.RHVAEs.var"#19#20"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils._finite_difference_gradient), Type{Vector{Float32}}, AutoEncoderToolkit.RHVAEs.var"#19#20"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [22] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.RHVAEs.var"#19#20"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [25] #∇hamiltonian#50
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:2076 [inlined]
   [26] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##∇hamiltonian#50", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Symbol, @NamedTuple{}, typeof(AutoEncoderToolkit.RHVAEs.∇hamiltonian), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, Symbol}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [27] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.∇hamiltonian), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, Symbol}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [28] #_leapfrog_third_step#56
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:2686 [inlined]
   [29] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##_leapfrog_third_step#56", Float32, typeof(AutoEncoderToolkit.RHVAEs.∇hamiltonian_TaylorDiff), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_third_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [30] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_third_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [31] #general_leapfrog_step#58
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:2907 [inlined]
   [32] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##general_leapfrog_step#58", Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.RHVAEs.general_leapfrog_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:ϵ, :∇H_kwargs)}}, Tuple{Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{ϵ::Float32, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Nothing, true}}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:ϵ, :steps, :∇H_kwargs)}}, Tuple{Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_second_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:ϵ, :steps, :∇H_kwargs)}}, Tuple{Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_first_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_third_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{JointGaussianLogDecoder{…}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(identity), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, typeof(identity)}}, Zygote.var"#1766#1767"{Zygote.var"#1764#1765"}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, JointGaussianLogDecoder{…}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, JointGaussianLogDecoder{…}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Zygote.Pullback{Tuple{Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:layers, Zygote.Context{false}, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}, Zygote.Pullback{Tuple{typeof(Flux._applychain), Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.var"#243#244"{Zygote.var"#back#back##0"{2, 1, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Zygote.Pullback{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(relu), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, typeof(relu)}}, Zygote.ZBack{NNlib.var"#broadcasted_relu_pullback#rrule##21"{Vector{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.Pullback{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(relu), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, typeof(relu)}}, Zygote.ZBack{NNlib.var"#broadcasted_relu_pullback#rrule##21"{Vector{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.var"#243#244"{Zygote.var"#back#back##0"{2, 2, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:μ, :logσ)}}, Tuple{Vector{Float32}, Vector{Float32}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, Nothing, true}}}}, Zygote.Pullback{Tuple{Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(identity), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, typeof(identity)}}, Zygote.var"#1766#1767"{Zygote.var"#1764#1765"}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, JointGaussianLogDecoder{…}, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}}, Zygote.var"#232#233"{typeof(identity)}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.slogdet), Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##slogdet#17", Bool, typeof(AutoEncoderToolkit.utils.slogdet), Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{check::Bool}, typeof(AutoEncoderToolkit.utils._slogdet), Type{Matrix{Float32}}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:check,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{check::Bool}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Matrix{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##1"{Int64, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.G_inv), Vector{Float32}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:λ, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Float32}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.G_inv), Vector{Float32}, Matrix{Float32}, Array{Float32, 3}, Float32, Float32}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs._G_inv), Type{Vector{Float32}}, Vector{Float32}, Matrix{Float32}, Array{Float32, 3}, Float32, Float32}, Tuple{Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcasted), typeof(exp), Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcastable), Matrix{Float32}}, Tuple{}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing, Nothing, Nothing}, Tuple{}}, Zygote.var"#1949#1950"{Zygote.var"#bc_fwd_back#_broadcast_forward##2"{Matrix{ForwardDiff.Dual{Nothing, Float32, 1}}, Tuple{Matrix{Float32}}, Val{1}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#890#891"{Zygote.var"#map_back#∇map##2"{typeof(Base.Broadcast.broadcastable), 1, Tuple{Tuple{}}, Tuple{Val{0}}, Tuple{}}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing}, Tuple{}}, Zygote.var"#combine_styles_pullback#combine_styles_pullback##0"{Tuple{Nothing, Nothing}}}}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.RHVAEs.var"#_G_inv##0#_G_inv##1"{Vector{Float32}, Float32}, Nothing, false}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#1522#1523"{Zygote.var"#1520#1521"}, Zygote.ZBack{Zygote.var"#literal_pow_pullback#literal_pow_pullback##0"{2, Float32}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.var"#kw_zpullback#chain_rrule_kw##0"{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Int64, Array{Float32, 3}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}}}}}, Zygote.var"#1548#1549"{Zygote.var"#1546#1547"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Float32}, Any}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Vector{Float32}, Matrix{Float32}}}, Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64, Int64, Colon}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#kw_zpullback#chain_rrule_kw##0"{ChainRules.var"#dropdims_pullback#384"{ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}}, Zygote.ZBack{ChainRules.var"#slash_pullback_scalar#rrule##224"{Matrix{Float32}, Float32}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Array{Float32, 3}}, Tuple{}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#1697#1698"{Zygote.var"#1695#1696"{Array{Float32, 3}, Array{Float32, 3}}}}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:M, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Array{Float32, 3}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:centroids_latent, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:T, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Float32}}}}}})(Δ::Tuple{Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [33] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, G_inv::typeof(AutoEncoderToolkit.RHVAEs.G_inv)}, typeof(AutoEncoderToolkit.RHVAEs.general_leapfrog_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}}, Any})(Δ::Tuple{Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 4 more times ---
   [38] #riemannian_hamiltonian_elbo#72
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:3878 [inlined]
   [39] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##riemannian_hamiltonian_elbo#72", Float32, Int64, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.riemannian_hamiltonian_elbo), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Vector{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [40] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, G_inv::typeof(AutoEncoderToolkit.RHVAEs.G_inv), tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), logp_prefactor::Vector{Float32}, logq_prefactor::Vector{Float32}}, typeof(AutoEncoderToolkit.RHVAEs.riemannian_hamiltonian_elbo), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Vector{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [42] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:-1 [inlined]
   [43] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##loss#76", Int64, Float32, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :position_logprior, :momentum_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, Nothing, true}}}}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [44] (::Zygote.Pullback{Tuple{var"#84#85", AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##loss#76", Int64, Float32, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :position_logprior, :momentum_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, Nothing, true}}}}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [45] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#84#85", AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##loss#76", Int64, Float32, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :position_logprior, :momentum_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, Nothing, true}}}}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [46] gradient(f::Function, args::AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [47] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1545
   [48] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [49] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1547 [inlined]
   [50] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [51] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1558 [inlined]
  --- the above 2 lines are repeated 1 more time ---
   [54] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [55] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [56] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [57] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:42 [inlined]
   [58] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [59] top-level scope
      @ none:6
   [60] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [61] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
gradient (different input and output): Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1562
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
    [4] unit_vector
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1058 [inlined]
    [5] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [6] #_unit_vectors##0
      @ ./none:-1 [inlined]
    [7] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] #∇map##7
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202 [inlined]
    [9] #Generator##0
      @ ./generator.jl:37 [inlined]
   [10] iterate
      @ ./generator.jl:48 [inlined]
   [11] collect(itr::Base.Generator{Base.Iterators.Zip{Tuple{Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Vector{Vector{Float32}}}}, Base.var"#Generator##0#Generator##1"{Zygote.var"#∇map##7#∇map##8"}})
      @ Base ./array.jl:790
   [12] map
      @ ./abstractarray.jl:3536 [inlined]
   [13] (::Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202
   [14] collect_pullback
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:236 [inlined]
   [15] _unit_vectors
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1127 [inlined]
   [16] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [17] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vectors), Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [20] #finite_difference_gradient#29
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1194 [inlined]
   [21] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##finite_difference_gradient#29", Symbol, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.RHVAEs.var"#19#20"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils._finite_difference_gradient), Type{Vector{Float32}}, AutoEncoderToolkit.RHVAEs.var"#19#20"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [22] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.RHVAEs.var"#19#20"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [25] #∇hamiltonian#50
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:2076 [inlined]
   [26] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##∇hamiltonian#50", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Symbol, @NamedTuple{}, typeof(AutoEncoderToolkit.RHVAEs.∇hamiltonian), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, Symbol}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [27] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.∇hamiltonian), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, Symbol}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [28] #_leapfrog_third_step#56
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:2686 [inlined]
   [29] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##_leapfrog_third_step#56", Float32, typeof(AutoEncoderToolkit.RHVAEs.∇hamiltonian_TaylorDiff), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_third_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [30] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_third_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [31] #general_leapfrog_step#58
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:2907 [inlined]
   [32] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##general_leapfrog_step#58", Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.RHVAEs.general_leapfrog_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:ϵ, :∇H_kwargs)}}, Tuple{Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{ϵ::Float32, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Nothing, true}}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:ϵ, :steps, :∇H_kwargs)}}, Tuple{Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_second_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:ϵ, :steps, :∇H_kwargs)}}, Tuple{Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_first_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_third_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{JointGaussianLogDecoder{…}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(identity), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, typeof(identity)}}, Zygote.var"#1766#1767"{Zygote.var"#1764#1765"}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, JointGaussianLogDecoder{…}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, JointGaussianLogDecoder{…}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Zygote.Pullback{Tuple{Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:layers, Zygote.Context{false}, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}, Zygote.Pullback{Tuple{typeof(Flux._applychain), Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.var"#243#244"{Zygote.var"#back#back##0"{2, 1, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Zygote.Pullback{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(relu), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, typeof(relu)}}, Zygote.ZBack{NNlib.var"#broadcasted_relu_pullback#rrule##21"{Vector{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.Pullback{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(relu), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, typeof(relu)}}, Zygote.ZBack{NNlib.var"#broadcasted_relu_pullback#rrule##21"{Vector{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.var"#243#244"{Zygote.var"#back#back##0"{2, 2, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:μ, :logσ)}}, Tuple{Vector{Float32}, Vector{Float32}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, Nothing, true}}}}, Zygote.Pullback{Tuple{Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(identity), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, typeof(identity)}}, Zygote.var"#1766#1767"{Zygote.var"#1764#1765"}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, JointGaussianLogDecoder{…}, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}}, Zygote.var"#232#233"{typeof(identity)}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.slogdet), Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##slogdet#17", Bool, typeof(AutoEncoderToolkit.utils.slogdet), Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{check::Bool}, typeof(AutoEncoderToolkit.utils._slogdet), Type{Matrix{Float32}}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:check,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{check::Bool}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Matrix{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##1"{Int64, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.G_inv), Vector{Float32}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:λ, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Float32}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.G_inv), Vector{Float32}, Matrix{Float32}, Array{Float32, 3}, Float32, Float32}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs._G_inv), Type{Vector{Float32}}, Vector{Float32}, Matrix{Float32}, Array{Float32, 3}, Float32, Float32}, Tuple{Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcasted), typeof(exp), Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcastable), Matrix{Float32}}, Tuple{}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing, Nothing, Nothing}, Tuple{}}, Zygote.var"#1949#1950"{Zygote.var"#bc_fwd_back#_broadcast_forward##2"{Matrix{ForwardDiff.Dual{Nothing, Float32, 1}}, Tuple{Matrix{Float32}}, Val{1}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#890#891"{Zygote.var"#map_back#∇map##2"{typeof(Base.Broadcast.broadcastable), 1, Tuple{Tuple{}}, Tuple{Val{0}}, Tuple{}}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing}, Tuple{}}, Zygote.var"#combine_styles_pullback#combine_styles_pullback##0"{Tuple{Nothing, Nothing}}}}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.RHVAEs.var"#_G_inv##0#_G_inv##1"{Vector{Float32}, Float32}, Nothing, false}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#1522#1523"{Zygote.var"#1520#1521"}, Zygote.ZBack{Zygote.var"#literal_pow_pullback#literal_pow_pullback##0"{2, Float32}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.var"#kw_zpullback#chain_rrule_kw##0"{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Int64, Array{Float32, 3}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}}}}}, Zygote.var"#1548#1549"{Zygote.var"#1546#1547"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Float32}, Any}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Vector{Float32}, Matrix{Float32}}}, Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64, Int64, Colon}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#kw_zpullback#chain_rrule_kw##0"{ChainRules.var"#dropdims_pullback#384"{ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}}, Zygote.ZBack{ChainRules.var"#slash_pullback_scalar#rrule##224"{Matrix{Float32}, Float32}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Array{Float32, 3}}, Tuple{}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#1697#1698"{Zygote.var"#1695#1696"{Array{Float32, 3}, Array{Float32, 3}}}}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:M, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Array{Float32, 3}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:centroids_latent, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:T, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Float32}}}}}})(Δ::Tuple{Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [33] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, G_inv::typeof(AutoEncoderToolkit.RHVAEs.G_inv)}, typeof(AutoEncoderToolkit.RHVAEs.general_leapfrog_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}}, Any})(Δ::Tuple{Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 4 more times ---
   [38] #riemannian_hamiltonian_elbo#74
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:4129 [inlined]
   [39] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##riemannian_hamiltonian_elbo#74", Float32, Int64, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.riemannian_hamiltonian_elbo), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Vector{Float32}, Vector{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [40] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, G_inv::typeof(AutoEncoderToolkit.RHVAEs.G_inv), tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), logp_prefactor::Vector{Float32}, logq_prefactor::Vector{Float32}}, typeof(AutoEncoderToolkit.RHVAEs.riemannian_hamiltonian_elbo), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Vector{Float32}, Vector{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [42] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:-1 [inlined]
   [43] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, Vector{Float32}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :position_logprior, :momentum_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##loss#77", Int64, Float32, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, Vector{Float32}, Vector{Float32}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [44] (::Zygote.Pullback{Tuple{var"#86#87", AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, Vector{Float32}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :position_logprior, :momentum_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##loss#77", Int64, Float32, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, Vector{Float32}, Vector{Float32}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, Vector{Float32}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [45] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#86#87", AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, Vector{Float32}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :position_logprior, :momentum_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##loss#77", Int64, Float32, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}}, Vector{Float32}, Vector{Float32}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, Vector{Float32}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [46] gradient(f::Function, args::AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianLogDecoder{…}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [47] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1545
   [48] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [49] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1547 [inlined]
   [50] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [51] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1563 [inlined]
  --- the above 2 lines are repeated 1 more time ---
   [54] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [55] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [56] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [57] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:42 [inlined]
   [58] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [59] top-level scope
      @ none:6
   [60] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [61] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
gradient (same input and output): Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1557
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
    [4] unit_vector
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1058 [inlined]
    [5] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [6] #_unit_vectors##0
      @ ./none:-1 [inlined]
    [7] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] #∇map##7
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202 [inlined]
    [9] #Generator##0
      @ ./generator.jl:37 [inlined]
   [10] iterate
      @ ./generator.jl:48 [inlined]
   [11] collect(itr::Base.Generator{Base.Iterators.Zip{Tuple{Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Vector{Vector{Float32}}}}, Base.var"#Generator##0#Generator##1"{Zygote.var"#∇map##7#∇map##8"}})
      @ Base ./array.jl:790
   [12] map
      @ ./abstractarray.jl:3536 [inlined]
   [13] (::Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202
   [14] collect_pullback
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:236 [inlined]
   [15] _unit_vectors
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1127 [inlined]
   [16] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [17] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vectors), Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [20] #finite_difference_gradient#29
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1194 [inlined]
   [21] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##finite_difference_gradient#29", Symbol, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.RHVAEs.var"#19#20"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils._finite_difference_gradient), Type{Vector{Float32}}, AutoEncoderToolkit.RHVAEs.var"#19#20"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [22] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.RHVAEs.var"#19#20"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [25] #∇hamiltonian#50
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:2076 [inlined]
   [26] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##∇hamiltonian#50", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Symbol, @NamedTuple{}, typeof(AutoEncoderToolkit.RHVAEs.∇hamiltonian), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, Symbol}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [27] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.∇hamiltonian), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, Symbol}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [28] #_leapfrog_third_step#56
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:2686 [inlined]
   [29] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##_leapfrog_third_step#56", Float32, typeof(AutoEncoderToolkit.RHVAEs.∇hamiltonian_TaylorDiff), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_third_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [30] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_third_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [31] #general_leapfrog_step#58
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:2907 [inlined]
   [32] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##general_leapfrog_step#58", Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.RHVAEs.general_leapfrog_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:ϵ, :∇H_kwargs)}}, Tuple{Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{ϵ::Float32, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Nothing, true}}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:ϵ, :steps, :∇H_kwargs)}}, Tuple{Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_second_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:ϵ, :steps, :∇H_kwargs)}}, Tuple{Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_first_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_third_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{SplitGaussianLogDecoder{…}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:μ, :logσ)}}, Tuple{Vector{Float32}, Vector{Float32}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, SplitGaussianLogDecoder{…}, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}}, Zygote.Pullback{Tuple{Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:layers, Zygote.Context{false}, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}, Zygote.Pullback{Tuple{typeof(Flux._applychain), Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.var"#243#244"{Zygote.var"#back#back##0"{2, 1, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Zygote.Pullback{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(relu), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, typeof(relu)}}, Zygote.ZBack{NNlib.var"#broadcasted_relu_pullback#rrule##21"{Vector{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.Pullback{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(relu), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, typeof(relu)}}, Zygote.ZBack{NNlib.var"#broadcasted_relu_pullback#rrule##21"{Vector{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.var"#243#244"{Zygote.var"#back#back##0"{2, 2, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}}}}, Zygote.Pullback{Tuple{Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:layers, Zygote.Context{false}, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}, Zygote.Pullback{Tuple{typeof(Flux._applychain), Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.var"#243#244"{Zygote.var"#back#back##0"{2, 1, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Zygote.Pullback{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(relu), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, typeof(relu)}}, Zygote.ZBack{NNlib.var"#broadcasted_relu_pullback#rrule##21"{Vector{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.Pullback{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(relu), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, typeof(relu)}}, Zygote.ZBack{NNlib.var"#broadcasted_relu_pullback#rrule##21"{Vector{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.var"#243#244"{Zygote.var"#back#back##0"{2, 2, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, SplitGaussianLogDecoder{…}, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.slogdet), Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##slogdet#17", Bool, typeof(AutoEncoderToolkit.utils.slogdet), Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{check::Bool}, typeof(AutoEncoderToolkit.utils._slogdet), Type{Matrix{Float32}}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:check,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{check::Bool}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Matrix{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##1"{Int64, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.G_inv), Vector{Float32}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:λ, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Float32}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.G_inv), Vector{Float32}, Matrix{Float32}, Array{Float32, 3}, Float32, Float32}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs._G_inv), Type{Vector{Float32}}, Vector{Float32}, Matrix{Float32}, Array{Float32, 3}, Float32, Float32}, Tuple{Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcasted), typeof(exp), Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcastable), Matrix{Float32}}, Tuple{}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing, Nothing, Nothing}, Tuple{}}, Zygote.var"#1949#1950"{Zygote.var"#bc_fwd_back#_broadcast_forward##2"{Matrix{ForwardDiff.Dual{Nothing, Float32, 1}}, Tuple{Matrix{Float32}}, Val{1}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#890#891"{Zygote.var"#map_back#∇map##2"{typeof(Base.Broadcast.broadcastable), 1, Tuple{Tuple{}}, Tuple{Val{0}}, Tuple{}}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing}, Tuple{}}, Zygote.var"#combine_styles_pullback#combine_styles_pullback##0"{Tuple{Nothing, Nothing}}}}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.RHVAEs.var"#_G_inv##0#_G_inv##1"{Vector{Float32}, Float32}, Nothing, false}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#1522#1523"{Zygote.var"#1520#1521"}, Zygote.ZBack{Zygote.var"#literal_pow_pullback#literal_pow_pullback##0"{2, Float32}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.var"#kw_zpullback#chain_rrule_kw##0"{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Int64, Array{Float32, 3}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}}}}}, Zygote.var"#1548#1549"{Zygote.var"#1546#1547"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Float32}, Any}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Vector{Float32}, Matrix{Float32}}}, Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64, Int64, Colon}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#kw_zpullback#chain_rrule_kw##0"{ChainRules.var"#dropdims_pullback#384"{ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}}, Zygote.ZBack{ChainRules.var"#slash_pullback_scalar#rrule##224"{Matrix{Float32}, Float32}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Array{Float32, 3}}, Tuple{}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#1697#1698"{Zygote.var"#1695#1696"{Array{Float32, 3}, Array{Float32, 3}}}}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:M, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Array{Float32, 3}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:centroids_latent, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:T, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Float32}}}}}})(Δ::Tuple{Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [33] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, G_inv::typeof(AutoEncoderToolkit.RHVAEs.G_inv)}, typeof(AutoEncoderToolkit.RHVAEs.general_leapfrog_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}}, Any})(Δ::Tuple{Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 4 more times ---
   [38] #riemannian_hamiltonian_elbo#72
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:3878 [inlined]
   [39] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##riemannian_hamiltonian_elbo#72", Float32, Int64, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.riemannian_hamiltonian_elbo), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Vector{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [40] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, G_inv::typeof(AutoEncoderToolkit.RHVAEs.G_inv), tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), logp_prefactor::Vector{Float32}, logq_prefactor::Vector{Float32}}, typeof(AutoEncoderToolkit.RHVAEs.riemannian_hamiltonian_elbo), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Vector{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [42] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:-1 [inlined]
   [43] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##loss#76", Int64, Float32, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :position_logprior, :momentum_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, Nothing, true}}}}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [44] (::Zygote.Pullback{Tuple{var"#84#85", AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##loss#76", Int64, Float32, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :position_logprior, :momentum_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, Nothing, true}}}}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [45] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#84#85", AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##loss#76", Int64, Float32, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :position_logprior, :momentum_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, Nothing, true}}}}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [46] gradient(f::Function, args::AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [47] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1545
   [48] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [49] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1547 [inlined]
   [50] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [51] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1558 [inlined]
  --- the above 2 lines are repeated 1 more time ---
   [54] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [55] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [56] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [57] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:42 [inlined]
   [58] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [59] top-level scope
      @ none:6
   [60] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [61] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
gradient (different input and output): Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1562
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
    [4] unit_vector
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1058 [inlined]
    [5] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [6] #_unit_vectors##0
      @ ./none:-1 [inlined]
    [7] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] #∇map##7
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202 [inlined]
    [9] #Generator##0
      @ ./generator.jl:37 [inlined]
   [10] iterate
      @ ./generator.jl:48 [inlined]
   [11] collect(itr::Base.Generator{Base.Iterators.Zip{Tuple{Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Vector{Vector{Float32}}}}, Base.var"#Generator##0#Generator##1"{Zygote.var"#∇map##7#∇map##8"}})
      @ Base ./array.jl:790
   [12] map
      @ ./abstractarray.jl:3536 [inlined]
   [13] (::Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202
   [14] collect_pullback
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:236 [inlined]
   [15] _unit_vectors
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1127 [inlined]
   [16] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [17] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vectors), Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [20] #finite_difference_gradient#29
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1194 [inlined]
   [21] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##finite_difference_gradient#29", Symbol, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.RHVAEs.var"#19#20"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils._finite_difference_gradient), Type{Vector{Float32}}, AutoEncoderToolkit.RHVAEs.var"#19#20"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [22] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.RHVAEs.var"#19#20"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Vector{Float32}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [25] #∇hamiltonian#50
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:2076 [inlined]
   [26] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##∇hamiltonian#50", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Symbol, @NamedTuple{}, typeof(AutoEncoderToolkit.RHVAEs.∇hamiltonian), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, Symbol}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [27] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.∇hamiltonian), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, Symbol}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [28] #_leapfrog_third_step#56
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:2686 [inlined]
   [29] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##_leapfrog_third_step#56", Float32, typeof(AutoEncoderToolkit.RHVAEs.∇hamiltonian_TaylorDiff), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_third_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [30] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_third_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [31] #general_leapfrog_step#58
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:2907 [inlined]
   [32] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##general_leapfrog_step#58", Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.RHVAEs.general_leapfrog_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:ϵ, :∇H_kwargs)}}, Tuple{Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{ϵ::Float32, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Nothing, true}}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:ϵ, :steps, :∇H_kwargs)}}, Tuple{Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_second_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:ϵ, :steps, :∇H_kwargs)}}, Tuple{Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_first_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_third_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{SplitGaussianLogDecoder{…}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:μ, :logσ)}}, Tuple{Vector{Float32}, Vector{Float32}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:logσ, Zygote.Context{false}, SplitGaussianLogDecoder{…}, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}}, Zygote.Pullback{Tuple{Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:layers, Zygote.Context{false}, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}, Zygote.Pullback{Tuple{typeof(Flux._applychain), Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.var"#243#244"{Zygote.var"#back#back##0"{2, 1, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Zygote.Pullback{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(relu), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, typeof(relu)}}, Zygote.ZBack{NNlib.var"#broadcasted_relu_pullback#rrule##21"{Vector{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.Pullback{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(relu), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, typeof(relu)}}, Zygote.ZBack{NNlib.var"#broadcasted_relu_pullback#rrule##21"{Vector{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.var"#243#244"{Zygote.var"#back#back##0"{2, 2, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}}}}, Zygote.Pullback{Tuple{Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:layers, Zygote.Context{false}, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}, Zygote.Pullback{Tuple{typeof(Flux._applychain), Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.var"#243#244"{Zygote.var"#back#back##0"{2, 1, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Zygote.Pullback{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(relu), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, typeof(relu)}}, Zygote.ZBack{NNlib.var"#broadcasted_relu_pullback#rrule##21"{Vector{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.Pullback{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(relu), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, typeof(relu)}}, Zygote.ZBack{NNlib.var"#broadcasted_relu_pullback#rrule##21"{Vector{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.var"#243#244"{Zygote.var"#back#back##0"{2, 2, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, SplitGaussianLogDecoder{…}, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.slogdet), Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##slogdet#17", Bool, typeof(AutoEncoderToolkit.utils.slogdet), Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{check::Bool}, typeof(AutoEncoderToolkit.utils._slogdet), Type{Matrix{Float32}}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:check,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{check::Bool}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Matrix{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##1"{Int64, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.G_inv), Vector{Float32}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:λ, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Float32}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.G_inv), Vector{Float32}, Matrix{Float32}, Array{Float32, 3}, Float32, Float32}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs._G_inv), Type{Vector{Float32}}, Vector{Float32}, Matrix{Float32}, Array{Float32, 3}, Float32, Float32}, Tuple{Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcasted), typeof(exp), Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcastable), Matrix{Float32}}, Tuple{}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing, Nothing, Nothing}, Tuple{}}, Zygote.var"#1949#1950"{Zygote.var"#bc_fwd_back#_broadcast_forward##2"{Matrix{ForwardDiff.Dual{Nothing, Float32, 1}}, Tuple{Matrix{Float32}}, Val{1}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#890#891"{Zygote.var"#map_back#∇map##2"{typeof(Base.Broadcast.broadcastable), 1, Tuple{Tuple{}}, Tuple{Val{0}}, Tuple{}}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing}, Tuple{}}, Zygote.var"#combine_styles_pullback#combine_styles_pullback##0"{Tuple{Nothing, Nothing}}}}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.RHVAEs.var"#_G_inv##0#_G_inv##1"{Vector{Float32}, Float32}, Nothing, false}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#1522#1523"{Zygote.var"#1520#1521"}, Zygote.ZBack{Zygote.var"#literal_pow_pullback#literal_pow_pullback##0"{2, Float32}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.var"#kw_zpullback#chain_rrule_kw##0"{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Int64, Array{Float32, 3}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}}}}}, Zygote.var"#1548#1549"{Zygote.var"#1546#1547"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Float32}, Any}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Vector{Float32}, Matrix{Float32}}}, Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64, Int64, Colon}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#kw_zpullback#chain_rrule_kw##0"{ChainRules.var"#dropdims_pullback#384"{ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}}, Zygote.ZBack{ChainRules.var"#slash_pullback_scalar#rrule##224"{Matrix{Float32}, Float32}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Array{Float32, 3}}, Tuple{}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#1697#1698"{Zygote.var"#1695#1696"{Array{Float32, 3}, Array{Float32, 3}}}}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:M, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Array{Float32, 3}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:centroids_latent, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:T, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Float32}}}}}})(Δ::Tuple{Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [33] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, G_inv::typeof(AutoEncoderToolkit.RHVAEs.G_inv)}, typeof(AutoEncoderToolkit.RHVAEs.general_leapfrog_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianLogDecoder{…}, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}}, Any})(Δ::Tuple{Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, @NamedTuple{μ::Vector{Float32}, logσ::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 4 more times ---
   [38] #riemannian_hamiltonian_elbo#74
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:4129 [inlined]
   [39] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##riemannian_hamiltonian_elbo#74", Float32, Int64, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.riemannian_hamiltonian_elbo), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Vector{Float32}, Vector{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [40] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, G_inv::typeof(AutoEncoderToolkit.RHVAEs.G_inv), tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), logp_prefactor::Vector{Float32}, logq_prefactor::Vector{Float32}}, typeof(AutoEncoderToolkit.RHVAEs.riemannian_hamiltonian_elbo), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Vector{Float32}, Vector{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [42] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:-1 [inlined]
   [43] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, Vector{Float32}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :position_logprior, :momentum_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##loss#77", Int64, Float32, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, Vector{Float32}, Vector{Float32}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [44] (::Zygote.Pullback{Tuple{var"#86#87", AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, Vector{Float32}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :position_logprior, :momentum_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##loss#77", Int64, Float32, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, Vector{Float32}, Vector{Float32}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, Vector{Float32}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [45] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#86#87", AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, Vector{Float32}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :position_logprior, :momentum_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##loss#77", Int64, Float32, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}}, Vector{Float32}, Vector{Float32}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, Vector{Float32}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [46] gradient(f::Function, args::AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, SplitGaussianLogDecoder{…}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [47] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1545
   [48] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [49] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1547 [inlined]
   [50] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [51] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1563 [inlined]
  --- the above 2 lines are repeated 1 more time ---
   [54] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [55] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [56] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [57] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:42 [inlined]
   [58] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [59] top-level scope
      @ none:6
   [60] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [61] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
gradient (same input and output): Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1557
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
    [4] unit_vector
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1058 [inlined]
    [5] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [6] #_unit_vectors##0
      @ ./none:-1 [inlined]
    [7] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] #∇map##7
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202 [inlined]
    [9] #Generator##0
      @ ./generator.jl:37 [inlined]
   [10] iterate
      @ ./generator.jl:48 [inlined]
   [11] collect(itr::Base.Generator{Base.Iterators.Zip{Tuple{Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Vector{Vector{Float32}}}}, Base.var"#Generator##0#Generator##1"{Zygote.var"#∇map##7#∇map##8"}})
      @ Base ./array.jl:790
   [12] map
      @ ./abstractarray.jl:3536 [inlined]
   [13] (::Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202
   [14] collect_pullback
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:236 [inlined]
   [15] _unit_vectors
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1127 [inlined]
   [16] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [17] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vectors), Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [20] #finite_difference_gradient#29
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1194 [inlined]
   [21] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##finite_difference_gradient#29", Symbol, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.RHVAEs.var"#19#20"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils._finite_difference_gradient), Type{Vector{Float32}}, AutoEncoderToolkit.RHVAEs.var"#19#20"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [22] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.RHVAEs.var"#19#20"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [25] #∇hamiltonian#50
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:2076 [inlined]
   [26] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##∇hamiltonian#50", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Symbol, @NamedTuple{}, typeof(AutoEncoderToolkit.RHVAEs.∇hamiltonian), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}, Symbol}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [27] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.∇hamiltonian), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}, Symbol}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [28] #_leapfrog_third_step#56
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:2686 [inlined]
   [29] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##_leapfrog_third_step#56", Float32, typeof(AutoEncoderToolkit.RHVAEs.∇hamiltonian_TaylorDiff), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_third_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [30] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_third_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [31] #general_leapfrog_step#58
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:2907 [inlined]
   [32] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##general_leapfrog_step#58", Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.RHVAEs.general_leapfrog_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:ϵ, :∇H_kwargs)}}, Tuple{Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{ϵ::Float32, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Nothing, true}}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:ϵ, :steps, :∇H_kwargs)}}, Tuple{Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_second_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:ϵ, :steps, :∇H_kwargs)}}, Tuple{Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_first_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_third_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{JointGaussianDecoder{…}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Dense{typeof(softplus), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(softplus), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(softplus), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(softplus), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(softplus), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(softplus), Matrix{Float32}, Vector{Float32}}, typeof(softplus)}}, Zygote.ZBack{NNlib.var"#broadcasted_softplus_pullback#rrule##81"{Vector{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, JointGaussianDecoder{…}, Dense{typeof(softplus), Matrix{Float32}, Vector{Float32}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, JointGaussianDecoder{…}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Zygote.Pullback{Tuple{Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:layers, Zygote.Context{false}, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}, Zygote.Pullback{Tuple{typeof(Flux._applychain), Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.var"#243#244"{Zygote.var"#back#back##0"{2, 1, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Zygote.Pullback{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(relu), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, typeof(relu)}}, Zygote.ZBack{NNlib.var"#broadcasted_relu_pullback#rrule##21"{Vector{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.Pullback{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(relu), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, typeof(relu)}}, Zygote.ZBack{NNlib.var"#broadcasted_relu_pullback#rrule##21"{Vector{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.var"#243#244"{Zygote.var"#back#back##0"{2, 2, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:μ, :σ)}}, Tuple{Vector{Float32}, Vector{Float32}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}, Nothing, true}}}}, Zygote.Pullback{Tuple{Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(identity), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, typeof(identity)}}, Zygote.var"#1766#1767"{Zygote.var"#1764#1765"}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, JointGaussianDecoder{…}, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}}, Zygote.var"#232#233"{typeof(identity)}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.slogdet), Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##slogdet#17", Bool, typeof(AutoEncoderToolkit.utils.slogdet), Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{check::Bool}, typeof(AutoEncoderToolkit.utils._slogdet), Type{Matrix{Float32}}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:check,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{check::Bool}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Matrix{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##1"{Int64, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.G_inv), Vector{Float32}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:λ, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Float32}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.G_inv), Vector{Float32}, Matrix{Float32}, Array{Float32, 3}, Float32, Float32}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs._G_inv), Type{Vector{Float32}}, Vector{Float32}, Matrix{Float32}, Array{Float32, 3}, Float32, Float32}, Tuple{Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcasted), typeof(exp), Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcastable), Matrix{Float32}}, Tuple{}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing, Nothing, Nothing}, Tuple{}}, Zygote.var"#1949#1950"{Zygote.var"#bc_fwd_back#_broadcast_forward##2"{Matrix{ForwardDiff.Dual{Nothing, Float32, 1}}, Tuple{Matrix{Float32}}, Val{1}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#890#891"{Zygote.var"#map_back#∇map##2"{typeof(Base.Broadcast.broadcastable), 1, Tuple{Tuple{}}, Tuple{Val{0}}, Tuple{}}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing}, Tuple{}}, Zygote.var"#combine_styles_pullback#combine_styles_pullback##0"{Tuple{Nothing, Nothing}}}}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.RHVAEs.var"#_G_inv##0#_G_inv##1"{Vector{Float32}, Float32}, Nothing, false}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#1522#1523"{Zygote.var"#1520#1521"}, Zygote.ZBack{Zygote.var"#literal_pow_pullback#literal_pow_pullback##0"{2, Float32}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.var"#kw_zpullback#chain_rrule_kw##0"{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Int64, Array{Float32, 3}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}}}}}, Zygote.var"#1548#1549"{Zygote.var"#1546#1547"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Float32}, Any}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Vector{Float32}, Matrix{Float32}}}, Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64, Int64, Colon}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#kw_zpullback#chain_rrule_kw##0"{ChainRules.var"#dropdims_pullback#384"{ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}}, Zygote.ZBack{ChainRules.var"#slash_pullback_scalar#rrule##224"{Matrix{Float32}, Float32}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Array{Float32, 3}}, Tuple{}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#1697#1698"{Zygote.var"#1695#1696"{Array{Float32, 3}, Array{Float32, 3}}}}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:M, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Array{Float32, 3}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:centroids_latent, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:T, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Float32}}}}}})(Δ::Tuple{Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [33] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, G_inv::typeof(AutoEncoderToolkit.RHVAEs.G_inv)}, typeof(AutoEncoderToolkit.RHVAEs.general_leapfrog_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}}, Any})(Δ::Tuple{Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 4 more times ---
   [38] #riemannian_hamiltonian_elbo#72
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:3878 [inlined]
   [39] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##riemannian_hamiltonian_elbo#72", Float32, Int64, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.riemannian_hamiltonian_elbo), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Vector{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [40] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, G_inv::typeof(AutoEncoderToolkit.RHVAEs.G_inv), tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), logp_prefactor::Vector{Float32}, logq_prefactor::Vector{Float32}}, typeof(AutoEncoderToolkit.RHVAEs.riemannian_hamiltonian_elbo), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Vector{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [42] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:-1 [inlined]
   [43] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##loss#76", Int64, Float32, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :position_logprior, :momentum_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, Nothing, true}}}}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [44] (::Zygote.Pullback{Tuple{var"#84#85", AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##loss#76", Int64, Float32, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :position_logprior, :momentum_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, Nothing, true}}}}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [45] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#84#85", AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##loss#76", Int64, Float32, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :position_logprior, :momentum_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, Nothing, true}}}}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [46] gradient(f::Function, args::AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [47] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1545
   [48] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [49] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1547 [inlined]
   [50] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [51] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1558 [inlined]
  --- the above 2 lines are repeated 1 more time ---
   [54] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [55] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [56] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [57] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:42 [inlined]
   [58] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [59] top-level scope
      @ none:6
   [60] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [61] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
gradient (different input and output): Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1562
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
    [4] unit_vector
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1058 [inlined]
    [5] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [6] #_unit_vectors##0
      @ ./none:-1 [inlined]
    [7] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] #∇map##7
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202 [inlined]
    [9] #Generator##0
      @ ./generator.jl:37 [inlined]
   [10] iterate
      @ ./generator.jl:48 [inlined]
   [11] collect(itr::Base.Generator{Base.Iterators.Zip{Tuple{Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Vector{Vector{Float32}}}}, Base.var"#Generator##0#Generator##1"{Zygote.var"#∇map##7#∇map##8"}})
      @ Base ./array.jl:790
   [12] map
      @ ./abstractarray.jl:3536 [inlined]
   [13] (::Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202
   [14] collect_pullback
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:236 [inlined]
   [15] _unit_vectors
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1127 [inlined]
   [16] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [17] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vectors), Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [20] #finite_difference_gradient#29
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1194 [inlined]
   [21] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##finite_difference_gradient#29", Symbol, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.RHVAEs.var"#19#20"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils._finite_difference_gradient), Type{Vector{Float32}}, AutoEncoderToolkit.RHVAEs.var"#19#20"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [22] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.RHVAEs.var"#19#20"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [25] #∇hamiltonian#50
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:2076 [inlined]
   [26] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##∇hamiltonian#50", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Symbol, @NamedTuple{}, typeof(AutoEncoderToolkit.RHVAEs.∇hamiltonian), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}, Symbol}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [27] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.∇hamiltonian), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}, Symbol}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [28] #_leapfrog_third_step#56
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:2686 [inlined]
   [29] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##_leapfrog_third_step#56", Float32, typeof(AutoEncoderToolkit.RHVAEs.∇hamiltonian_TaylorDiff), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_third_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [30] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_third_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [31] #general_leapfrog_step#58
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:2907 [inlined]
   [32] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##general_leapfrog_step#58", Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.RHVAEs.general_leapfrog_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:ϵ, :∇H_kwargs)}}, Tuple{Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{ϵ::Float32, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Nothing, true}}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:ϵ, :steps, :∇H_kwargs)}}, Tuple{Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_second_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:ϵ, :steps, :∇H_kwargs)}}, Tuple{Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_first_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_third_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{JointGaussianDecoder{…}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Dense{typeof(softplus), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(softplus), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(softplus), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(softplus), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(softplus), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(softplus), Matrix{Float32}, Vector{Float32}}, typeof(softplus)}}, Zygote.ZBack{NNlib.var"#broadcasted_softplus_pullback#rrule##81"{Vector{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, JointGaussianDecoder{…}, Dense{typeof(softplus), Matrix{Float32}, Vector{Float32}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:μ, Zygote.Context{false}, JointGaussianDecoder{…}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Zygote.Pullback{Tuple{Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:layers, Zygote.Context{false}, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}, Zygote.Pullback{Tuple{typeof(Flux._applychain), Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.var"#243#244"{Zygote.var"#back#back##0"{2, 1, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Zygote.Pullback{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(relu), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, typeof(relu)}}, Zygote.ZBack{NNlib.var"#broadcasted_relu_pullback#rrule##21"{Vector{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.Pullback{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(relu), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, typeof(relu)}}, Zygote.ZBack{NNlib.var"#broadcasted_relu_pullback#rrule##21"{Vector{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.var"#243#244"{Zygote.var"#back#back##0"{2, 2, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:μ, :σ)}}, Tuple{Vector{Float32}, Vector{Float32}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}, Nothing, true}}}}, Zygote.Pullback{Tuple{Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(identity), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, typeof(identity)}}, Zygote.var"#1766#1767"{Zygote.var"#1764#1765"}, Zygote.ZBack{ChainRules.var"#times_pullback#rrule##129"{Matrix{Float32}, Vector{Float32}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Any}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:decoder, Zygote.Context{false}, JointGaussianDecoder{…}, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}}, Zygote.var"#232#233"{typeof(identity)}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.slogdet), Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##slogdet#17", Bool, typeof(AutoEncoderToolkit.utils.slogdet), Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{check::Bool}, typeof(AutoEncoderToolkit.utils._slogdet), Type{Matrix{Float32}}, Matrix{Float32}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:check,)}}, Tuple{Bool}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{check::Bool}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Matrix{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#-_pullback##1"{Int64, ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.G_inv), Vector{Float32}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:λ, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Float32}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.G_inv), Vector{Float32}, Matrix{Float32}, Array{Float32, 3}, Float32, Float32}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs._G_inv), Type{Vector{Float32}}, Vector{Float32}, Matrix{Float32}, Array{Float32, 3}, Float32, Float32}, Tuple{Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcasted), typeof(exp), Matrix{Float32}}, Tuple{Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcastable), Matrix{Float32}}, Tuple{}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing, Nothing, Nothing}, Tuple{}}, Zygote.var"#1949#1950"{Zygote.var"#bc_fwd_back#_broadcast_forward##2"{Matrix{ForwardDiff.Dual{Nothing, Float32, 1}}, Tuple{Matrix{Float32}}, Val{1}}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#890#891"{Zygote.var"#map_back#∇map##2"{typeof(Base.Broadcast.broadcastable), 1, Tuple{Tuple{}}, Tuple{Val{0}}, Tuple{}}}, Zygote.var"#361#362"{Zygote.var"#359#360"{Tuple{Tuple{Nothing}, Tuple{}}, Zygote.var"#combine_styles_pullback#combine_styles_pullback##0"{Tuple{Nothing, Nothing}}}}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.RHVAEs.var"#_G_inv##0#_G_inv##1"{Vector{Float32}, Float32}, Nothing, false}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#1522#1523"{Zygote.var"#1520#1521"}, Zygote.ZBack{Zygote.var"#literal_pow_pullback#literal_pow_pullback##0"{2, Float32}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.var"#kw_zpullback#chain_rrule_kw##0"{ChainRules.var"#mean_pullback#696"{Int64, ChainRules.var"#sum_pullback#550"{Int64, Array{Float32, 3}, ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}}}}}, Zygote.var"#1548#1549"{Zygote.var"#1546#1547"}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Float32}, Any}, Zygote.var"#1671#1672"{Zygote.var"#1669#1670"{Vector{Float32}, Matrix{Float32}}}, Zygote.var"#805#806"{Zygote.var"#801#802"{Matrix{Float32}, Tuple{Int64, Int64, Colon}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#kw_zpullback#chain_rrule_kw##0"{ChainRules.var"#dropdims_pullback#384"{ChainRulesCore.ProjectTo{AbstractArray, @NamedTuple{element::ChainRulesCore.ProjectTo{Float32, @NamedTuple{}}, axes::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}}, Zygote.ZBack{ChainRules.var"#slash_pullback_scalar#rrule##224"{Matrix{Float32}, Float32}}, Zygote.var"#1035#1036"{Zygote.var"#1029#1030"{Matrix{Float32}}}, Zygote.var"#162#163"{Zygote.var"#160#161"}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#1753#1754"{Zygote.var"#1751#1752"{2, Matrix{Float32}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Array{Float32, 3}}, Tuple{}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:dims,)}}, Tuple{Int64}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{dims::Int64}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Matrix{Float32}}, Tuple{}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#1697#1698"{Zygote.var"#1695#1696"{Array{Float32, 3}, Array{Float32, 3}}}}}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:M, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Array{Float32, 3}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:centroids_latent, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:T, Zygote.Context{false}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Float32}}}}}})(Δ::Tuple{Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [33] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, G_inv::typeof(AutoEncoderToolkit.RHVAEs.G_inv)}, typeof(AutoEncoderToolkit.RHVAEs.general_leapfrog_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, JointGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}}, Any})(Δ::Tuple{Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 4 more times ---
   [38] #riemannian_hamiltonian_elbo#74
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:4129 [inlined]
   [39] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##riemannian_hamiltonian_elbo#74", Float32, Int64, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Bool, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.riemannian_hamiltonian_elbo), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Vector{Float32}, Vector{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [40] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{K::Int64, ϵ::Float32, βₒ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, G_inv::typeof(AutoEncoderToolkit.RHVAEs.G_inv), tempering_schedule::typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), logp_prefactor::Vector{Float32}, logq_prefactor::Vector{Float32}}, typeof(AutoEncoderToolkit.RHVAEs.riemannian_hamiltonian_elbo), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}, Vector{Float32}, Vector{Float32}}, Any})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
   [42] loss
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:-1 [inlined]
   [43] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, Vector{Float32}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :position_logprior, :momentum_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##loss#77", Int64, Float32, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, Vector{Float32}, Vector{Float32}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [44] (::Zygote.Pullback{Tuple{var"#86#87", AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, Vector{Float32}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :position_logprior, :momentum_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##loss#77", Int64, Float32, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, Vector{Float32}, Vector{Float32}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, Vector{Float32}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [45] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{var"#86#87", AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}}, Tuple{Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglobal), Module, Symbol}, Tuple{typeof(getglobal)}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, Vector{Float32}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.ZBack{Zygote.var"#Real_pullback#Real_pullback##0"}, Zygote.Pullback{Tuple{Type{NamedTuple{(:reconstruction_loglikelihood, :position_logprior, :momentum_logprior)}}, Tuple{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, Nothing, true}}}}, Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##loss#77", Int64, Float32, Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.HVAEs.quadratic_tempering), Nothing, @NamedTuple{}, Float32, Vector{Float32}, Vector{Float32}, typeof(AutoEncoderToolkit.RHVAEs.loss), AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}}, Vector{Float32}, Vector{Float32}}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}}}, Zygote.var"#215#216"{Zygote.var"#213#214"{Zygote.Context{false}, GlobalRef, Vector{Float32}}}}}})(Δ::Float32)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:91
   [46] gradient(f::Function, args::AutoEncoderToolkit.RHVAEs.RHVAE{AutoEncoderToolkit.VAEs.VAE{JointGaussianLogEncoder{…}, JointGaussianDecoder{…}}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface.jl:148
   [47] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1545
   [48] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [49] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1547 [inlined]
   [50] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [51] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1563 [inlined]
  --- the above 2 lines are repeated 1 more time ---
   [54] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [55] top-level scope
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:18
   [56] macro expansion
      @ /opt/julia/share/julia/stdlib/v1.12/Test/src/Test.jl:1724 [inlined]
   [57] macro expansion
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/test/runtests.jl:42 [inlined]
   [58] include(mapexpr::Function, mod::Module, _path::String)
      @ Base ./Base.jl:301
   [59] top-level scope
      @ none:6
   [60] eval(m::Module, e::Any)
      @ Core ./boot.jl:485
   [61] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:295
gradient (same input and output): Error During Test at /home/pkgeval/.julia/packages/AutoEncoderToolkit/YfAJl/test/rhvae.jl:1557
  Got exception outside of a @test
  Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_has_free_typevars), Int32, svec(Any), 0, :(:ccall), %2)).
  You might want to check the Zygote limitations documentation.
  https://fluxml.ai/Zygote.jl/latest/limitations
  
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:44
    [2] (::Zygote.Pullback{Tuple{typeof(Core.has_free_typevars), Type{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Zygote.var"#IntX_pullback#IntX_pullback##1"}}})(Δ::Nothing)
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 1 more time ---
    [4] unit_vector
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1058 [inlined]
    [5] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [6] #_unit_vectors##0
      @ ./none:-1 [inlined]
    [7] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
    [8] #∇map##7
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202 [inlined]
    [9] #Generator##0
      @ ./generator.jl:37 [inlined]
   [10] iterate
      @ ./generator.jl:48 [inlined]
   [11] collect(itr::Base.Generator{Base.Iterators.Zip{Tuple{Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}, Vector{Vector{Float32}}}}, Base.var"#Generator##0#Generator##1"{Zygote.var"#∇map##7#∇map##8"}})
      @ Base ./array.jl:790
   [12] map
      @ ./abstractarray.jl:3536 [inlined]
   [13] (::Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:202
   [14] collect_pullback
      @ ~/.julia/packages/Zygote/TWpme/src/lib/array.jl:236 [inlined]
   [15] _unit_vectors
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1127 [inlined]
   [16] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [17] (::Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vectors), Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils._unit_vectors), Type{Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}}, Nothing, false}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Vector{Float32}}, Any}, Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Vector{Float32}, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Int64}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:x, Zygote.Context{false}, AutoEncoderToolkit.utils.var"#_unit_vectors##0#_unit_vectors##1"{Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.unit_vector), Vector{Float32}, Int64}, Tuple{Zygote.var"#collect_pullback#_pullback##9"{Zygote.var"#map_back#∇map##2"{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, 1, Tuple{UnitRange{Int64}}, Tuple{Tuple{Base.OneTo{Int64}}}, Vector{Tuple{Float32, Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Int64}, Any}}}}, Nothing}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Int64}, Any}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Core._typeof_captured_variable), Type{Float32}}, Any}, Zygote.var"#389#390"{Zygote.Jnew{AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, Nothing, false}}, Zygote.Pullback{Tuple{Type{Base.Generator}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}, UnitRange{Int64}}, Tuple{Zygote.var"#389#390"{Zygote.Jnew{Base.Generator{UnitRange{Int64}, AutoEncoderToolkit.utils.var"#unit_vector##0#unit_vector##1"{Int64, Type{Float32}}}, Nothing, false}}}}, Zygote.Pullback{Tuple{typeof(eltype), Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}, Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}}}}}}}}, Nothing}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}}})(Δ::Vector{Vector{Float32}})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [20] #finite_difference_gradient#29
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/utils.jl:1194 [inlined]
   [21] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.utils.var"##finite_difference_gradient#29", Symbol, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.RHVAEs.var"#19#20"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:fdtype,)}}, Tuple{Symbol}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{fdtype::Symbol}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(AutoEncoderToolkit.utils.storage_type), Vector{Float32}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils._finite_difference_gradient), Type{Vector{Float32}}, AutoEncoderToolkit.RHVAEs.var"#19#20"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Any}}})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [22] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{fdtype::Symbol}, typeof(AutoEncoderToolkit.utils.finite_difference_gradient), AutoEncoderToolkit.RHVAEs.var"#19#20"{typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Vector{Float32}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
  --- the above 1 lines are repeated 2 more times ---
   [25] #∇hamiltonian#50
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:2076 [inlined]
   [26] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##∇hamiltonian#50", typeof(AutoEncoderToolkit.decoder_loglikelihood), typeof(AutoEncoderToolkit.spherical_logprior), typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior), Symbol, @NamedTuple{}, typeof(AutoEncoderToolkit.RHVAEs.∇hamiltonian), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}, Symbol}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [27] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.∇hamiltonian), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}, Symbol}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [28] #_leapfrog_third_step#56
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:2686 [inlined]
   [29] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##_leapfrog_third_step#56", Float32, typeof(AutoEncoderToolkit.RHVAEs.∇hamiltonian_TaylorDiff), @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_third_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:0
   [30] (::Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_third_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any})(Δ::Vector{Float32})
      @ Zygote ~/.julia/packages/Zygote/TWpme/src/compiler/interface2.jl:110
   [31] #general_leapfrog_step#58
      @ ~/.julia/packages/AutoEncoderToolkit/YfAJl/src/rhvae.jl:2907 [inlined]
   [32] (::Zygote.Pullback{Tuple{AutoEncoderToolkit.RHVAEs.var"##general_leapfrog_step#58", Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}, typeof(AutoEncoderToolkit.RHVAEs.G_inv), typeof(AutoEncoderToolkit.RHVAEs.general_leapfrog_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}, @NamedTuple{centroids_latent::Matrix{Float32}, L::Array{Float32, 3}, M::Array{Float32, 3}, T::Float32, λ::Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:ϵ, :∇H_kwargs)}}, Tuple{Float32, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{ϵ::Float32, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Nothing, true}}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:ϵ, :steps, :∇H_kwargs)}}, Tuple{Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_second_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple{(:ϵ, :steps, :∇H_kwargs)}}, Tuple{Float32, Int64, @NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, Nothing, true}}}}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, steps::Int64, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_first_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{typeof(Core.kwcall), @NamedTuple{ϵ::Float32, ∇H_kwargs::@NamedTuple{reconstruction_loglikelihood::typeof(AutoEncoderToolkit.decoder_loglikelihood), position_logprior::typeof(AutoEncoderToolkit.spherical_logprior), momentum_logprior::typeof(AutoEncoderToolkit.RHVAEs.riemannian_logprior)}}, typeof(AutoEncoderToolkit.RHVAEs._leapfrog_third_step), Vector{Float32}, Vector{Float32}, Vector{Float32}, Matrix{Float32}, Float32, SplitGaussianDecoder{…}, @NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}}, Any}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#232#233"{typeof(identity)}, Zygote.Pullback{Tuple{SplitGaussianDecoder{…}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:μ, :σ)}}, Tuple{Vector{Float32}, Vector{Float32}}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#393#394"{Zygote.Jnew{@NamedTuple{μ::Vector{Float32}, σ::Vector{Float32}}, Nothing, true}}}}, Zygote.var"#232#233"{typeof(identity)}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:σ, Zygote.Context{false}, SplitGaussianDecoder{…}, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(softplus), Matrix{Float32}, Vector{Float32}}}}}}, Zygote.Pullback{Tuple{Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(softplus), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}}, Tuple{Zygote.var"#373#374"{Zygote.var"#back#back##9"{:layers, Zygote.Context{false}, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(softplus), Matrix{Float32}, Vector{Float32}}}}, Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(softplus), Matrix{Float32}, Vector{Float32}}}}}, Zygote.Pullback{Tuple{typeof(Flux._applychain), Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(softplus), Matrix{Float32}, Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.var"#243#244"{Zygote.var"#back#back##0"{2, 1, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Zygote.Pullback{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.ZBack{Returns{Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#1658#1659"{Zygote.var"#1654#1655"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.ZBack{Flux.var"#206#207"}, Zygote.ZBack{Returns{NTuple{4, ChainRulesCore.NoTangent}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(relu), Vector{Float32}}, Tuple{}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:weight, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.var"#373#374"{Zygote.var"#back#back##9"{:bias, Zygote.Context{false}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(getproperty), Module, Symbol}, Tuple{Zygote.Pullback{Tuple{typeof(getglob
PkgEval terminated after 2282.0s: test log exceeded the size limit