Block-representation of dense tensors

[jofrevalles]

I tested Tensor with BlockArray in data field and the results work as expected:

julia> using BlockArrays
julia> using Random
julia> Random.seed!(123)
TaskLocalRNG()

julia> block_data_a = [rand(4, 4) for _ in 1:4, _ in 1:4]
4×4 Matrix{Matrix{Float64}}:
 [0.9063 0.253849 0.0991336 0.0320967...]

julia> block_data_b = [rand(4, 4) for _ in 1:4, _ in 1:4]
4×4 Matrix{Matrix{Float64}}:
 [0.255027 0.868985 0.95932 0.823697...]

julia> # Create the block arrays
       block_array_a = BlockArray(block_data_a, [2, 2], [2, 2])
2×2-blocked 4×4 BlockMatrix{Matrix{Float64}}:
 [0.9063 0.253849 0.0991336 0.0320967; 0.443494 0.334152 0.125287 0.350546; 0.745673 0.427328 0.692209 0.930332; 0.512083 0.867547 0.136551 0.959434]   …  [0.732574 0.972201 0.877295 0.146259; 0.385228 0.0610019 0.893793 0.433955; 0.798934 0.980256 0.337006 0.291577; 0.829194 0.225583 0.251766 0.827852]
 [0.919181 0.954159 0.789493 0.123538; 0.426019 0.845895 0.619259 0.74002; 0.746586 0.586749 0.477645 0.705747; 0.819201 0.121813 0.804193 0.991961]  
 ─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────     [0.122918 0.906864 0.13252 0.519136; 0.0432075 0.91711 0.831918 0.0766932; 0.260182 0.537621 0.915793 0.877763; 0.888027 0.220951 0.95494 0.882158]  
 ─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
 [0.675195 ... 0.52969]    

julia> block_array_b = BlockArray(block_data_b, [2, 2], [2, 2])
2×2-blocked 4×4 BlockMatrix{Matrix{Float64}}:
 [0.255027 0.868985 0.95932 0.823697; 0.597402 0.522017 0.646037 0.185613; 0.729799 0.767848 0.694181 0.813438; 0.398507 0.500034 0.832089 0.621681]   …  [0.262472 0.205768 0.500511 0.334961; 0.472211 0.920132 0.671532 0.186199; 0.128476 0.399426 0.869231 0.830299; 0.05036 0.52651 0.678405 0.930738]   
 [0.224888 0.971681 0.441624 0.736363; 0.41986 0.138707 0.570679 0.350911; 0.2386 0.771563 0.684492 0.739064; 0.327655 0.263582 0.563876 0.244367]   
 ────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────     [0.0556571 0.515958 0.837037 0.704039; 0.686477 0.674318 0.91214 0.0217106; 0.819139 0.907658 0.668988 0.737698; 0.747796 0.308451 0.131705 0.772859]
 ─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
 [0.533516 0.0114595 0.806676 0.785328...] 

julia> tensor_a = Tensor(block_array_a, (:a, :b))
4×4 Tensor{Matrix{Float64}, 2, BlockMatrix{Matrix{Float64}, Matrix{Matrix{Matrix{Float64}}}, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}}} with indices 1:1:4×1:1:4: ...

julia> tensor_b = Tensor(block_array_b, (:b, :c))
4×4 Tensor{Matrix{Float64}, 2, BlockMatrix{Matrix{Float64}, Matrix{Matrix{Matrix{Float64}}}, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}}} with indices 1:1:4×1:1:4: ...

julia> # Perform the contraction
       result = contract(tensor_a, tensor_b)
4×4 Tensor{Matrix{Float64}, 2, BlockMatrix{Matrix{Float64}, Matrix{Matrix{Matrix{Float64}}}, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}}} with indices 1:1:4×1:1:4:
 [2.70393 2.44856 4.53006 4.51218; 2.38024 2.37521 4.27813 4.00492; 3.731 4.42418 5.94166 5.72008; 3.75654 3.23434 6.50645 5.04544]    …  [2.55026 2.92253 3.44155 2.99779; 2.76131 3.53482 3.21373 3.78507; 2.33087 4.97109 5.44422 5.1327; 3.57959 5.49615 5.20459 5.22641]
 [3.81849 3.89957 5.66889 4.97532; 3.71879 4.37482 5.69095 5.15109; 3.75261 4.52058 6.39526 6.04347; 3.60494 4.03185 6.54128 5.58345]     [3.12488 4.2088 5.11801 3.77261; 3.49575 4.15328 4.58628 4.72482; 3.78557 5.23826 5.23467 5.59499; 3.33741 4.59042 4.75271 6.0863]
 [3.27923 4.05041 5.5721 5.0491; 3.04664 3.48437 4.5289 4.28216; 4.27774 4.72328 7.06008 6.45713; 2.27613 2.60933 4.08217 3.44402]        [3.23618 4.26344 4.88191 4.63693; 2.59837 3.91299 4.3737 4.13523; 3.78282 5.64347 5.77367 6.09598; 2.42089 3.18015 3.2306 3.35643]
 [2.32643 3.252 3.95212 3.91159; 3.09044 3.86356 5.29118 4.96522; 4.06865 4.88741 7.0316 6.49199; 3.90975 4.34506 6.54487 6.25266]        [2.25425 4.15514 4.02621 3.95229; 3.13965 4.29838 4.25514 4.62259; 3.58519 5.15794 5.6374 6.25417; 4.11128 5.84972 5.76475 5.84386]

[mofeing]

Great! Could you write some integration tests where you try contract, permutedims, ...?

Next steps: which factorizations are implemented? Does eigendecomposition work? SVD? QR?

[jofrevalles]

I have created integration tests for BlockArray, covering functions such as contract, permutedims, and basic manipulation of the tensors. You can find the tests in this Pull Request.

Unfortunately, I found that svd decomposition is currently not supported for BlockArrays. Here's a relevant GitHub issue discussing the lack of SVD support in BlockArrays.jl: https://github.com/JuliaArrays/BlockArrays.jl/issues/131 I will look at other block-array libraries to see if they have such implementations.

[mofeing]

Unfortunately, I found that svd decomposition is currently not supported for BlockArrays. Here's a relevant GitHub issue discussing the lack of SVD support in BlockArrays.jl: JuliaArrays/BlockArrays.jl#131 I will look at other block-array libraries to see if they have such implementations.

Okay, this is what I expected. Don't worry, distributed contraction has higher priority right now.

[mofeing]

Have you tested eigendecomposition? And QR?

[jofrevalles]

Have you tested eigendecomposition? And QR?

Not yet since we don't have these implementations in Tensors right now, but if you want I can try to implement these and test the integrations.

[mofeing]

Nah, just try it with a BlockMatrix.

[jofrevalles]

I tried it, and these two are not yet implemented as well. I added the tests commented in the PR.

[jofrevalles]

I've been evaluating the Dagger.jl library, specifically its Dagger.DArray implementation, in the context of our tensor computations. My findings have revealed several significant limitations:

1. Lack of support for the permutedims function: While the transpose function works as expected with DArray, the permutedims function does not. This issue is not documented in the library's repository, which suggests it may be worth raising as a new issue for the developers to address.

   
   julia> darray = distribute(rand(4, 4), Blocks(2, 2))
   DArray{Any, 2, typeof(cat)}(4, 4)

julia> transpose(darray) # works fine Dagger.Transpose{Any, 2}(4, 4)

julia> transpose(darray) |> compute |> collect 4×4 Matrix{Float64}: 0.167282 0.111796 0.139091 0.76079 0.970018 0.962297 0.0806159 0.750641 0.44158 0.213501 0.929178 0.0804493 0.980127 0.0202904 0.873157 0.249399

julia> permutedims(darray, (2, 1)) 4×4 Matrix{Any}: Dagger.GetIndex{Any, 2}(1, 0, 1, 0) Dagger.GetIndex{Any, 2}(1, 0, 1, 0) Dagger.GetIndex{Any, 2}(1, 0, 1, 0) Dagger.GetIndex{Any, 2}(1, 0, 1, 0) Dagger.GetIndex{Any, 2}(1, 0, 1, 0) Dagger.GetIndex{Any, 2}(1, 0, 1, 0) Dagger.GetIndex{Any, 2}(1, 0, 1, 0) Dagger.GetIndex{Any, 2}(1, 0, 1, 0) Dagger.GetIndex{Any, 2}(1, 0, 1, 0) Dagger.GetIndex{Any, 2}(1, 0, 1, 0) Dagger.GetIndex{Any, 2}(1, 0, 1, 0) Dagger.GetIndex{Any, 2}(1, 0, 1, 0) Dagger.GetIndex{Any, 2}(1, 0, 1, 0) Dagger.GetIndex{Any, 2}(1, 0, 1, 0) Dagger.GetIndex{Any, 2}(1, 0, 1, 0) Dagger.GetIndex{Any, 2}(1, 0, 1, 0)

julia> permutedims(darray, (2, 1)) |> compute ERROR: MethodError: no method matching compute(::Context, ::Matrix{Any}; options=nothing) Closest candidates are: compute(::Context, ::DArray; persist, options) at ~/.julia/packages/Dagger/vNUsP/src/array/darray.jl:231 compute(::Context, ::Thunk; options) at ~/.julia/packages/Dagger/vNUsP/src/compute.jl:27 compute(::Any; options) at ~/.julia/packages/Dagger/vNUsP/src/compute.jl:5 ...

2. Incompatibility between `Dagger.DArray` and regular arrays: The library does not support operations between regular arrays and `Dagger.DArray`. For instance, attempting to multiply a `DArray `with a regular array results in an error:
```julia
julia> darray = distribute(rand(4, 4), Blocks(2, 2))
DArray{Any, 2, typeof(cat)}(4, 4)

julia> array = rand(4, 4)
4×4 Matrix{Float64}:
 0.503586  0.981538  0.631702  0.722101
 0.370623  0.485568  0.618103  0.239512
 0.69361   0.9219    0.68611   0.238639
 0.825931  0.29894   0.685966  0.0812979

julia> darray * array
ERROR: MethodError: no method matching *(::Dagger.Chunk{Float64, MemPool.DRef, Dagger.ThreadProc, AnyScope}, ::Float64)
Closest candidates are:
  *(::Any, ::Any, ::Any, ::Any...) at operators.jl:591
  *(::T, ::T) where T<:Union{Float16, Float32, Float64} at float.jl:385
  *(::StridedArray{P}, ::Real) where P<:Dates.Period at /opt/julia-1.8.2/share/julia/stdlib/v1.8/Dates/src/deprecated.jl:44
  ...

3. No support for LinearAlgebra functions: Similar to BlockArray, Dagger.DArray does not support key LinearAlgebra functions such as qr, svd, and eigen.

In summary, although BlockArrays seems a more developed option compared to Dagger.DArray, the task-oriented computational model of Dagger has the potential to provide superior capabilities for our parallel and distributed computations. @mofeing Do you want me to write an integration test for this one?

[mofeing]

@jofrevalles maybe this message should go in #7?

[jofrevalles]

@jofrevalles maybe this message should go in #7?

Well, here we use Dagger.DArray which has a block representation. But sure these two issues are related.