Combining LinearAlgebra.Diagonal with a CuArray inside @tensor

[OliverDudgeon]

When using normal arrays, I can use the Diagonal lazy matrix constructor and @tensor works normally.

A = Diagonal(rand(4))
B = rand(4,4)

# this works
@tensor out[a,c] := B[a,b] * A[b,c]

However, when using CuArrays the macro fails despite * working for these types.

Ag = CuArray(A)
Bg = CuArray(B)

Bg * Ag # this works
# but this doesn't
@tensor out[a,c] := Bg[a,b] * Ag[b,c]

We get an error message from Strided MethodError: no method matching StridedViews.StridedView(::LinearAlgebra.Diagonal{Float64, CUDA.CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}).

Is there any way to get this working?

[lkdvos]

I think it depends a bit on what you would like. The current implementation uses cuTENSOR to implement arbitrary contractions between CuArrays, and as far as I know this does not support Diagonal. Of course, in the specific case of matrix multiplication, we could use CUDA to compute the result. I would think it is possible to play some stride tricks to get cuTENSOR to work with diagonal matrices in the cases where the contractions are not just matrix multiplication but I can't say I currently have the time to play around with it. However, if you want to give it a go, I could point you to some examples of what I have in mind

[OliverDudgeon]

Thanks for the explanation. In my case, I was able to do the scaling due to the Diagonal matrix separately to the @tensor macro, so this isn't too important to me now.