update adapt_transformer

Author: lkdvos

ext/TensorKitAMDGPUExt/roctensormap.jl

@@ -163,11 +163,4 @@ for f in (:sqrt, :log, :asin, :acos, :acosh, :atanh, :acoth)
     end
 end
 
-function TensorKit.adapt_transformer(
-        t::TensorKit.GenericTreeTransformer, data::ROCVector
-    )
-    new_data = map(t.data) do (U, structs_dst, structs_src)
-        return AMDGPU.Adapt.adapt(ROCArray, U), structs_dst, structs_src
-    end
-    return TensorKit.GenericTreeTransformer(new_data)
-end
+TensorKit.adapt_transformer(U::AbstractMatrix, ::Type{A}) where {A <: ROCVector} = AMDGPU.Adapt.adapt(ROCArray, U)

ext/TensorKitCUDAExt/cutensormap.jl

@@ -169,11 +169,4 @@ for f in (:sqrt, :log, :asin, :acos, :acosh, :atanh, :acoth)
     end
 end
 
-function TensorKit.adapt_transformer(
-        t::TensorKit.GenericTreeTransformer, data::CuVector
-    )
-    new_data = map(t.data) do (U, structs_dst, structs_src)
-        return CUDA.Adapt.adapt(CuArray, U), structs_dst, structs_src
-    end
-    return TensorKit.GenericTreeTransformer(new_data)
-end
+TensorKit.adapt_transformer(U::AbstractMatrix, ::Type{A}) where {A <: CuVector} = CUDA.Adapt.adapt(CuArray, U)

src/tensors/indexmanipulations.jl

@@ -624,7 +624,8 @@ function add_transform_kernel!(
                         ptriv, false, One(), Zero(), backend, allocator
                     )
                 end
-                mul!(buffer_dst, buffer_src, transpose(StridedView(U)))
+                U′ = adapt_transformer(U, storagetype(tdst))
+                mul!(buffer_dst, buffer_src, transpose(StridedView(U′)))
                 @inbounds for (i, (f₃, f₄)) in enumerate(fusiontrees(dst))
                     TO.tensoradd!(
                         tdst[f₃, f₄], sreshape(buffer_dst[:, i], sz_src),
@@ -661,7 +662,6 @@ function add_transform_kernel!(
         data_dst::DenseVector, data_src::DenseVector, p, transformer::GenericTreeTransformer,
         α, β, backend, allocator, scheduler
     )
-    transformer = adapt_transformer(transformer, data_dst)
     # Each entry covers one fusion block:
     #   U            — recoupling matrix (rows = dst trees, cols = src trees)
     #   sz_{dst,src} — array shape of each block (same for all trees in the block)
@@ -697,7 +697,8 @@ function add_transform_kernel!(
 
             # 2. Recoupling: buffer_dst = buffer_src * U^T  (each output tree is a linear
             #    combination of input trees weighted by the recoupling coefficients).
-            mul!(buffer_dst, buffer_src, transpose(StridedView(U)))
+            U′ = adapt_transformer(U, typeof(data_dst))
+            mul!(buffer_dst, buffer_src, transpose(StridedView(U′)))
 
             # 3. Insert: scatter column i of buffer_dst into the destination, applying the
             #    actual index permutation p in the same tensoradd! call.
@@ -713,3 +714,12 @@ function add_transform_kernel!(
     TO.allocator_reset!(allocator, cp)
     return nothing
 end
+
+"""
+    adapt_transformer(U::AbstractMatrix, ::Type{A})
+
+Return a version of the basis transformation `U` that is compatible for storage type `A`.
+Default is a no-op.
+Backends (e.g. CUDA, AMDGPU) should overload this for their vector types to ensure the recoupling matrix `U` is on the correct device.
+"""
+adapt_transformer(U::AbstractMatrix, ::Type{A}) where {A} = U

src/tensors/treetransformers.jl

@@ -203,12 +203,3 @@ end
 function _transformer_weight((mat, structs_dst, structs_src)::GenericTransformerData)
     return length(mat) * prod(structs_dst[1])
 end
-
-"""
-    adapt_transformer(transformer::TreeTransformer, data::AbstractVector)
-
-Return a version of `transformer` whose internal arrays are compatible with `data`.
-Default is a no-op. Backends (e.g. CUDA, AMDGPU) should overload this for their vector types
-to ensure the recoupling matrix `U` inside `GenericTreeTransformer` is on the correct device.
-"""
-adapt_transformer(t::TreeTransformer, ::AbstractVector) = t