BlockSparseArrays.jl/src/abstractblocksparsearray/broadcast.jl at 37c0a86ccd31dc19b8e849303f239594fa67a487 · ITensor/BlockSparseArrays.jl · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
using BlockArrays: AbstractBlockedUnitRange, BlockSlice
using Base.Broadcast: BroadcastStyle

function Base.Broadcast.BroadcastStyle(arraytype::Type{<:AnyAbstractBlockSparseArray})
    return Broadcast.BlockSparseArrayStyle(BroadcastStyle(blocktype(arraytype)))
end

# Fix ambiguity error with `BlockArrays`.
function Base.Broadcast.BroadcastStyle(
        arraytype::Type{
            <:SubArray{
                <:Any,
                <:Any,
                <:AbstractBlockSparseArray,
                <:Tuple{BlockSlice{<:Any, <:Any, <:AbstractBlockedUnitRange}, Vararg{Any}},
            },
        },
    )
    return Broadcast.BlockSparseArrayStyle{ndims(arraytype)}()
end
function Base.Broadcast.BroadcastStyle(
        arraytype::Type{
            <:SubArray{
                <:Any,
                <:Any,
                <:AbstractBlockSparseArray,
                <:Tuple{
                    BlockSlice{<:Any, <:Any, <:AbstractBlockedUnitRange},
                    BlockSlice{<:Any, <:Any, <:AbstractBlockedUnitRange},
                    Vararg{Any},
                },
            },
        },
    )
    return Broadcast.BlockSparseArrayStyle{ndims(arraytype)}()
end
function Base.Broadcast.BroadcastStyle(
        arraytype::Type{
            <:SubArray{
                <:Any,
                <:Any,
                <:AbstractBlockSparseArray,
                <:Tuple{Any, BlockSlice{<:Any, <:Any, <:AbstractBlockedUnitRange}, Vararg{Any}},
            },
        },
    )
    return Broadcast.BlockSparseArrayStyle{ndims(arraytype)}()
end

# These catch cases that aren't caught by the standard
# `BlockSparseArrayStyle` definition, and also fix
# ambiguity issues.
function Base.copyto!(dest::AnyAbstractBlockSparseArray, bc::Broadcasted)
    return copyto!_blocksparse(dest, bc)
end
function Base.copyto!(
        dest::AnyAbstractBlockSparseArray, bc::Broadcasted{<:Base.Broadcast.AbstractArrayStyle{0}}
    )
    return copyto!_blocksparse(dest, bc)
end
function Base.copyto!(
        dest::AnyAbstractBlockSparseArray{<:Any, N}, bc::Broadcasted{<:Broadcast.BlockSparseArrayStyle{N}}
    ) where {N}
    return copyto!_blocksparse(dest, bc)
end