diff --git a/ext/PEPSKitGPUArraysExt.jl b/ext/PEPSKitGPUArraysExt.jl index 13510fd5c..e76b0a3f2 100644 --- a/ext/PEPSKitGPUArraysExt.jl +++ b/ext/PEPSKitGPUArraysExt.jl @@ -12,4 +12,16 @@ function PEPSKit._column_vectors(::Type{A}, m::AbstractMatrix) where {A <: AnyGP end PEPSKit._indexable_diagview(m::AnyGPUArray) = collect(MatrixAlgebraKit.diagview(m)) +# Reduce into a 0-dimensional device array instead of returning a host scalar. `sdiag_pow` only +# feeds this into a broadcast, and a 0-dim array broadcasts as a scalar, so the value never has to +# come back to the CPU. Returning a scalar here would force a device sync on every call, and +# `sdiag_pow` runs once per bond per weight absorption in simple update. +function PEPSKit._maxabs(data::AnyGPUArray) + T = real(eltype(data)) + acc = similar(data, T, ()) + fill!(acc, zero(T)) + Base.mapreducedim!(abs, max, acc, data) + return acc +end + end diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index 4280c30af..a8039169a 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -17,7 +17,7 @@ using MatrixAlgebraKit: svd_pullback!, svd_trunc_pullback! using TensorKit using TensorKit: AdjointTensorMap, SectorDict using TensorKit: throw_invalid_innerproduct, similarstoragetype -using TensorKit.Factorizations: TruncationSpace, _notrunc_ind +using TensorKit.Factorizations: TruncationSpace using KrylovKit using KrylovKit: Lanczos, BlockLanczos diff --git a/src/algorithms/time_evolution/trotter_gate.jl b/src/algorithms/time_evolution/trotter_gate.jl index 37997e6b4..9503ca4d2 100644 --- a/src/algorithms/time_evolution/trotter_gate.jl +++ b/src/algorithms/time_evolution/trotter_gate.jl @@ -22,7 +22,7 @@ function gate_to_mpo( # convert to Vidal gauge _cluster_truncate!(Os, fill(notrunc(), N - 1)) # evenly distribute the (Inf) norm - nrms = _maxabs.(Os) + nrms = norm.(Os, Inf) fac = prod(nrms)^(1 / N) scale!.(Os, fac ./ nrms) # remove trivial legs in first/last tensor, and restore MPO convention diff --git a/src/utility/util.jl b/src/utility/util.jl index df25d2ae5..663c4ae50 100644 --- a/src/utility/util.jl +++ b/src/utility/util.jl @@ -53,7 +53,6 @@ function ChainRulesCore.rrule( ::typeof(sdiag_pow), s::AbstractTensorMap, pow::Real; tol::Real = eps(real(scalartype(s)))^(3 / 4), ) - tol *= _maxabs(s) spow = sdiag_pow(s, pow; tol) spow_minus1_conj = scale!(sdiag_pow(s', pow - 1; tol), pow) function sdiag_pow_pullback(c̄_)