# Using TensorOperations.jl when contracting tensors not following strict Einstein summation convention

**URL:** <https://discourse.julialang.org/t/using-tensoroperations-jl-when-contracting-tensors-not-following-strict-einstein-summation-convention/113034>\
**Category:** Specific Domains\
**Tags:** question, tensors, tensoroperations\
**Created:** [April 16, 2024, 6:55pm UTC](https://discourse.julialang.org/t/using-tensoroperations-jl-when-contracting-tensors-not-following-strict-einstein-summation-convention/113034 "2024-04-16T18:55:14Z")\
**Posts on this page:** 9\
**Page:** 1

<div class="post-metadata">

**Author:** ![Gattu\_Mytraya](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gattu_mytraya/32/45429_2.png) [@Gattu\_Mytraya](https://discourse.julialang.org/u/Gattu_Mytraya)\
**Post date:** [April 16, 2024, 6:55pm UTC](https://discourse.julialang.org/t/using-tensoroperations-jl-when-contracting-tensors-not-following-strict-einstein-summation-convention/113034/1 "2024-04-16T18:55:14Z")

</div>

I am trying to implement the following computation on a GPU.

```julia
@tensor S[i,j] = A[a1,j]*F[a1,i,a2]*B[a2,j]

```

As per the strict Einstein summation convention, this is not a valid expression as the index j appears twice on the right-hand side and, as such, does not work when used with TensorOperations.jl.  
At the moment, I create a larger intermediate array,

```julia
@tensor temporary_location[i,j,jprime] = A[a1,j]*F[a1,i,a2]*B[a2,jprime]

```

and then copy the diagonal part (i.e., temporary\_location[i,j,j]) to the array S. I would like to know if there is a way to avoid performing the unnecessary computations involved in the above workaround. The range of indices a1, i, and j are of the same order (up to a few hundred), and the range of the index a2 is much smaller than a1, i, or j.

I tried using Tullio.jl, but it seems about 20-25x times slower than TensorOperations.jl.

Any help would be appreciated.

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [April 16, 2024, 7:22pm UTC](https://discourse.julialang.org/t/using-tensoroperations-jl-when-contracting-tensors-not-following-strict-einstein-summation-convention/113034/2 "2024-04-16T19:22:39Z")

</div>

How does this compare to what you have now?

```julia
using TensorCast, TensorOperations
@cast tmp[a1, a2, j] := A[a1, j] * B[a2, j]
@tensor S[i, j] := F[a1, i, a2] * tmp[a1, a2, j]

```

---

<div class="post-metadata">

**Author:** ![juthohaegeman](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/juthohaegeman/32/8620_2.png) [@juthohaegeman](https://discourse.julialang.org/u/juthohaegeman)\
**Post date:** [April 16, 2024, 7:26pm UTC](https://discourse.julialang.org/t/using-tensoroperations-jl-when-contracting-tensors-not-following-strict-einstein-summation-convention/113034/3 "2024-04-16T19:26:57Z")

</div>

For maximal performance, it is probably better to just call the relevant CUBLAS functions for this contraction, as it is not particularly difficult.

```julia
temp = reshape(F, sizea1 * sizei, sizea2) * B
temp2 = reshape(temp, sizea1, sizei, sizej)

```

The final step is then a batched matrix vector multiplication, if you read it as  
`S[i, j] = temp2[a1, i, j] * A[a1, j]`. Here I found the function  
`gemv_strided_batched!`, so I think something like this should work:

```julia
gemv_strided_batched!('T', 1, temp2, A, 0., S)

```

---

<div class="post-metadata">

**Author:** ![Gattu\_Mytraya](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gattu_mytraya/32/45429_2.png) [@Gattu\_Mytraya](https://discourse.julialang.org/u/Gattu_Mytraya)\
**Post date:** [April 16, 2024, 8:00pm UTC](https://discourse.julialang.org/t/using-tensoroperations-jl-when-contracting-tensors-not-following-strict-einstein-summation-convention/113034/4 "2024-04-16T20:00:40Z")

</div>

There seems to be some issue with @cast.  
As per your suggestion, I use

```julia
@cast tmp[a1,a2,j] := A[a1,j]*B[a2,j]

```

I use arrays A of shape (28,90) and B of shape (9,90), but the output has shape (180,9,180) instead of shape (28,9,90). I am unsure what’s wrong.

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [April 16, 2024, 8:15pm UTC](https://discourse.julialang.org/t/using-tensoroperations-jl-when-contracting-tensors-not-following-strict-einstein-summation-convention/113034/5 "2024-04-16T20:15:59Z")

</div>

That should work, and does for me, I don’t know how you got that size. Here’s working code for all three versions, although not tried on a GPU:

```julia
A = rand(28,90);
B = rand(9,90);
F = rand(28,100,9);
using TensorOperations, Tullio, TensorCast, NNlib, TensorCore

@tullio S1[i,j] := A[a1,j]*F[a1,i,a2]*B[a2,j]; # original expression
size(S1) # (100, 90)
# CPU: min 167.541 μs, mean 173.655 μs (49 allocations, 73.17 KiB)

@cast tmp[a1, a2, j] := A[a1, j] * B[a2, j];
size(tmp) # (28, 9, 90)
@tensor S2[i, j] := F[a1, i, a2] * tmp[a1, a2, j];
S1 ≈ S2
# CPU: min 57.500 μs, mean 99.750 μs (16 allocations, 445.23 KiB)

temp2 = boxdot(F, B);
size(temp2) # (28, 100, 90)
temp3 = batched_mul(batched_transpose(temp2), reshape(A, size(A,1), 1, :));
S3 = reshape(temp3, size(temp3,1), :);
S1 ≈ S3
# CPU: min 301.583 μs, mean 485.143 μs (47 allocations, 2.00 MiB)

```

---

<div class="post-metadata">

**Author:** ![Gattu\_Mytraya](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gattu_mytraya/32/45429_2.png) [@Gattu\_Mytraya](https://discourse.julialang.org/u/Gattu_Mytraya)\
**Post date:** [April 16, 2024, 11:47pm UTC](https://discourse.julialang.org/t/using-tensoroperations-jl-when-contracting-tensors-not-following-strict-einstein-summation-convention/113034/6 "2024-04-16T23:47:03Z")

</div>

My bad. I used a different matrix than the one I intended to use.  
With

```julia
using TensorCast, TensorOperations
@cast tmp[a1, a2, j] := A[a1, j] * B[a2, j]
@tensor S[i, j] := F[a1, i, a2] * tmp[a1, a2, j]

```

I see a 2 times increase in speed in the case where dim(a1) = 28, dim(a2) = 9, dim(i)=dim(j)=90. Thanks.

---

<div class="post-metadata">

**Author:** ![Gattu\_Mytraya](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gattu_mytraya/32/45429_2.png) [@Gattu\_Mytraya](https://discourse.julialang.org/u/Gattu_Mytraya)\
**Post date:** [April 17, 2024, 12:08am UTC](https://discourse.julialang.org/t/using-tensoroperations-jl-when-contracting-tensors-not-following-strict-einstein-summation-convention/113034/7 "2024-04-17T00:08:27Z")

</div>

When I try your suggestion, I get the following error:

```julia
ERROR: CUBLASError: an invalid value was used as an argument (code 7, CUBLAS_STATUS_INVALID_VALUE)
Stacktrace:
 [1] throw_api_error(res::CUDA.CUBLAS.cublasStatus_t)
   @ CUDA.CUBLAS ~/work/.julia/packages/CUDA/rXson/lib/cublas/libcublas.jl:11
 [2] check
   @ ~/work/.julia/packages/CUDA/rXson/lib/cublas/libcublas.jl:21 [inlined]
 [3] cublasZgemvStridedBatched
   @ ~/work/.julia/packages/CUDA/rXson/lib/utils/call.jl:26 [inlined]
 [4] gemv_strided_batched!(trans::Char, alpha::Int64, A::CuArray{…}, x::CuArray{…}, beta::Float64, y::CuArray{…})
   @ CUDA.CUBLAS ~/work/.julia/packages/CUDA/rXson/lib/cublas/wrappers.jl:420

```

I am not sure what’s wrong. 'T" is for transpose, right? (In case it makes any difference, although, looking at the definition of gemv\_strided\_batched! it shouldn’t, my matrices are complex)

---

<div class="post-metadata">

**Author:** ![juthohaegeman](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/juthohaegeman/32/8620_2.png) [@juthohaegeman](https://discourse.julialang.org/u/juthohaegeman)\
**Post date:** [April 17, 2024, 8:23am UTC](https://discourse.julialang.org/t/using-tensoroperations-jl-when-contracting-tensors-not-following-strict-einstein-summation-convention/113034/8 "2024-04-17T08:23:17Z")

</div>

Can you post a fully functional minimal code (minimal working example, or in this case, not working but error reproducing)? For the contraction you want to do, it should indeed be transpose, even if the data is complex.

---

<div class="post-metadata">

**Author:** ![Gattu\_Mytraya](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gattu_mytraya/32/45429_2.png) [@Gattu\_Mytraya](https://discourse.julialang.org/u/Gattu_Mytraya)\
**Post date:** [April 17, 2024, 12:47pm UTC](https://discourse.julialang.org/t/using-tensoroperations-jl-when-contracting-tensors-not-following-strict-einstein-summation-convention/113034/9 "2024-04-17T12:47:46Z")

</div>

I am unsure what information this can convey, so do let me know if you need more.  
This is the code I am trying to execute:

```julia
    display(temp2)
    display(exp_θ_gpu)
    display(slater_det_gpu)
    CUDA.CUBLAS.gemv_strided_batched!('T', 1, temp2, exp_θ_gpu, 0., slater_det_gpu)

```

which gives me

```julia
3×4×4 CuArray{ComplexF64, 3, CUDA.Mem.DeviceBuffer}:
[:, :, 1] =
 0.0+0.0im 0.155265+0.170144im 0.24062-0.219578im -0.155265-0.170144im
 1.0+0.0im -0.31053+0.689596im -9.20336e-33-4.14433e-33im -0.31053+0.689596im
 0.0+0.0im 0.155265-0.85974im 1.21586+0.219578im -0.155265+0.85974im

[:, :, 2] =
 0.0+0.0im -0.0157624+0.595112im 0.841615+0.0222914im 0.0157624-0.595112im
 1.0+0.0im 0.0315249-0.160339im 2.13989e-33+4.20731e-34im 0.0315249-0.160339im
 0.0+0.0im -0.0157624-0.434772im 0.614861-0.0222914im 0.0157624+0.434772im

[:, :, 3] =
 0.0+0.0im -0.0675479+0.811143im 1.14713+0.0955272im 0.0675479-0.811143im
 1.0+0.0im 0.135096-0.592402im 7.9062e-33+1.80299e-33im 0.135096-0.592402im
 0.0+0.0im -0.0675479-0.218741im 0.309346-0.0955272im 0.0675479+0.218741im

[:, :, 4] =
 0.0+0.0im 0.131172+0.479609im 0.67827-0.185505im -0.131172-0.479609im
 1.0+0.0im -0.262344+0.070665im -9.43096e-34-3.50125e-33im -0.262344+0.070665im
 0.0+0.0im 0.131172-0.550274im 0.778205+0.185505im -0.131172+0.550274im
3×4 CuArray{ComplexF64, 2, CUDA.Mem.DeviceBuffer}:
 0.159602+0.987181im 0.658391+0.752676im -0.134971+0.99085im -0.924661+0.380791im
      1.0-0.0im 1.0-0.0im 1.0-0.0im 1.0-0.0im
 0.159602-0.987181im 0.658391-0.752676im -0.134971-0.99085im -0.924661-0.380791im
4×4 CuArray{ComplexF64, 2, CUDA.Mem.DeviceBuffer}:
 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im
 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im
 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im
 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im
ERROR: CUBLASError: an invalid value was used as an argument (code 7, CUBLAS_STATUS_INVALID_VALUE)
Stacktrace:
 [1] throw_api_error(res::CUDA.CUBLAS.cublasStatus_t)
   @ CUDA.CUBLAS ~/work/.julia/packages/CUDA/rXson/lib/cublas/libcublas.jl:11
 [2] check
   @ ~/work/.julia/packages/CUDA/rXson/lib/cublas/libcublas.jl:21 [inlined]
 [3] cublasZgemvStridedBatched
   @ ~/work/.julia/packages/CUDA/rXson/lib/utils/call.jl:26 [inlined]
 [4] gemv_strided_batched!(trans::Char, alpha::Int64, A::CuArray{…}, x::CuArray{…}, beta::Float64, y::CuArray{…})
   @ CUDA.CUBLAS ~/work/.julia/packages/CUDA/rXson/lib/cublas/wrappers.jl:420

```