add hook for adapt_transformer

Author: lkdvos

ext/TensorKitAMDGPUExt/roctensormap.jl

@@ -162,3 +162,12 @@ for f in (:sqrt, :log, :asin, :acos, :acosh, :atanh, :acoth)
         return tf
     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

ext/TensorKitCUDAExt/cutensormap.jl

@@ -169,6 +169,11 @@ for f in (:sqrt, :log, :asin, :acos, :acosh, :atanh, :acoth)
     end
 end
 
-function TensorKit._add_transform_multi!(tdst::CuTensorMap, tsrc, p, (U, structs_dst, structs_src)::Tuple{<:Array, TD, TS}, buffers, alpha, beta, backend...) where {TD, TS}
-    return TensorKit._add_transform_multi!(tdst, tsrc, p, (CUDA.Adapt.adapt(CuArray, U), structs_dst, structs_src), buffers, alpha, beta, backend...)
+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

src/tensors/indexmanipulations.jl

@@ -661,6 +661,7 @@ 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)

src/tensors/treetransformers.jl

@@ -203,3 +203,12 @@ 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