fix tensorcast issue

stecrotti · Apr 15, 2024 · 2a5147d · 2a5147d
1 parent bbc8d5b
commit 2a5147d
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/UniformTensorTrains/UniformTensorTrains.jl b/src/UniformTensorTrains/UniformTensorTrains.jl
@@ -5,6 +5,7 @@ using ..TensorTrains: _reshape1
 
 using LinearAlgebra
 using KrylovKit: eigsolve
+using TensorCast: TensorCast, @cast
 using Tullio: @tullio
 
 export AbstractPeriodicTensorTrain, PeriodicTensorTrain, flat_periodic_tt, rand_periodic_tt,

diff --git a/src/UniformTensorTrains/transfer_operator.jl b/src/UniformTensorTrains/transfer_operator.jl
@@ -48,8 +48,7 @@ end
 
 function leading_eig(G::AbstractTransferOperator)
     GG = collect(G)
-    # @cast B[(i,j),(k,l)] := GG[i,j,k,l]
-    B = reshape(GG, prod(size(GG)[1:2]), prod(size(GG)[3:4]))
+    @cast B[(i,j),(k,l)] := GG[i,j,k,l]
     valsR, vecsR = eigsolve(B)
     valsL, vecsL = eigsolve(B')
     valsR[1] ≈ valsL[1] || @warn "Leading eigenvalue for A and Aᵀ not equal, got $(valsR[1]) and $(valsL[1])"