Small fixes to let BlockTensorMaps work with GPU arrays

Author: kshyatt

src/linalg/factorizations.jl

@@ -1,5 +1,5 @@
 using MatrixAlgebraKit
-using MatrixAlgebraKit: AbstractAlgorithm, YALAPACK.BlasMat, Algorithm
+using MatrixAlgebraKit: AbstractAlgorithm, YALAPACK.BlasMat, Algorithm, diagview
 import MatrixAlgebraKit as MAK
 
 # Type piracy for defining the MAK rules on BlockArrays!
@@ -9,9 +9,8 @@ const BlockBlasMat{T <: MAK.BlasFloat} = BlockMatrix{T}
 
 function MatrixAlgebraKit.one!(A::BlockBlasMat)
     _one, _zero = one(eltype(A)), zero(eltype(A))
-    @inbounds for j in axes(A, 2), i in axes(A, 1)
-        A[i, j] = ifelse(i == j, _one, _zero)
-    end
+    A .= _zero
+    diagview(A) .= _one
     return A
 end
 

src/tensors/blocktensor.jl

@@ -123,10 +123,10 @@ end
 # ------------------------
 for (fname, felt) in ((:zeros, :zero), (:ones, :one))
     @eval begin
-        function Base.$fname(::Type{T}, V::TensorMapSumSpace) where {T}
-            TT = blocktensormaptype(spacetype(V), numout(V), numin(V), T)
+        function Base.$fname(::Type{TorA}, V::TensorMapSumSpace) where {TorA}
+            TT = blocktensormaptype(spacetype(V), numout(V), numin(V), TorA)
             t = TT(undef, V)
-            fill!(t, $felt(T))
+            fill!(t, $felt(scalartype(t)))
             return t
         end
     end
@@ -136,9 +136,9 @@ for randfun in (:rand, :randn, :randexp)
     randfun! = Symbol(randfun, :!)
     @eval begin
         function Random.$randfun(
-                rng::Random.AbstractRNG, ::Type{T}, V::TensorMapSumSpace
-            ) where {T}
-            TT = blocktensormaptype(spacetype(V), numout(V), numin(V), T)
+                rng::Random.AbstractRNG, ::Type{TorA}, V::TensorMapSumSpace
+            ) where {TorA}
+            TT = blocktensormaptype(spacetype(V), numout(V), numin(V), TorA)
             t = TT(undef, V)
             Random.$randfun!(rng, t)
             return t