correct some mask issues, improve bcast support

Author: rayegun

src/abstractgbarray.jl

@@ -54,12 +54,13 @@ Base.promote_rule(::Type{<:AbstractGBVector{T, F}}, ::Type{<:AbstractGBVector{T2
     GBVector{promote_type(T, T2), promote_type(F, F2)}
 
 function gbpromote_strip(A, B)
-    if A isa Transpose{<:Any, <:AbstractGBVector} && B isa Transpose{<:Any, <:AbstractGBVector}
-        return GBMatrix
-    else
-        return promote_type(strip_parameters(typeof(parent(A))), strip_parameters(typeof(parent(B))))
-    end
+    return promote_type(strip_parameters(typeof(parent(A))), strip_parameters(typeof(parent(B))))
 end
+gbpromote_strip(::Transpose{<:Any, <:AbstractGBVector}, ::Transpose{<:Any, <:AbstractGBVector}) = GBMatrix
+gbpromote_strip(::AbstractGBVector, ::Transpose{<:Any, <:AbstractGBVector}) = GBMatrix
+gbpromote_strip(::Transpose{<:Any, <:AbstractGBVector}, ::AbstractGBVector) = GBMatrix
+
+
 
 Base.IndexStyle(::AbstractGBArray) = IndexCartesian()
 Base.eltype(::AbstractGBArray{T, F}) where {T, F} = Union{T, F}
@@ -534,6 +535,7 @@ function subassign!(
     J, nj = idx(J)
     desc = _handledescriptor(desc; out=C, in1=A)
     mask = _handlemask!(desc, mask)
+    mask isa AbstractVector && length(I) == 1 && (mask = copy(mask'))
     I = decrement!(I)
     I !== J && (J = decrement!(J))
     rereshape = false
@@ -568,6 +570,7 @@ function subassign!(C::AbstractGBArray{T}, x, I, J;
     I !== J && (J = decrement!(J))
     desc = _handledescriptor(desc; out=C)
     mask = _handlemask!(desc, mask)
+    mask isa AbstractVector && length(I) == 1 && (mask = copy(mask'))
     _subassign(C, x, I, ni, J, nj, mask, _handleaccum(accum, storedeltype(C)), desc)
     increment!(I)
     I !== J && (decrement!(J))
@@ -637,6 +640,7 @@ function assign!(
     J, nj = idx(J)
     desc = _handledescriptor(desc; in1=A, out=C)
     mask = _handlemask!(desc, mask)
+    mask isa AbstractVector && length(I) == 1 && (mask = copy(mask'))
     I = decrement!(I)
     I !== J && (J = decrement!(J))
     if !(eltype(A) <: valid_union) || !(eltype(C) <: valid_union)
@@ -659,6 +663,7 @@ function assign!(C::AbstractGBArray{T}, x, I, J;
     I !== J && (J = decrement!(J))
     desc = _handledescriptor(desc; out=C)
     mask = _handlemask!(desc, mask)
+    mask isa AbstractVector && length(I) == 1 && (mask = copy(mask'))
     _assign(C, x, I, ni, J, nj, mask, _handleaccum(accum, storedeltype(C)), desc)
     increment!(I)
     I !== J && (decrement!(J))

src/indexutils.jl

@@ -34,18 +34,46 @@ function idx(I)
     end
 end
 
-# This function assumes that szA and szB are
-# technically equal and that
-# 1 <= length(szA | szB) <= 2
-# size checks should be done elsewhere.
-function _combinesizes(A, B)
-    if A isa Transpose{<:Any, <:AbstractVector} && B isa Transpose{<:Any, <:AbstractVector}
-        return size(A)
-    end
-    if (A isa AbstractVector && B isa AbstractMatrix) ||
-        (B isa AbstractVector && A isa AbstractMatrix)
-        return (size(A, 1), size(A, 2))
-    else
-        return size(A)
-    end
+# Combine sizes for bcasting purposes
+# This is quite inelegant :(, does this already exist somewhere?
+function _combinesizes(A::AbstractGBVector, B::AbstractGBVector)
+    size(A) == size(B) && (return size(A)) # same size
+    size(A, 1) == 1 && (return size(B)) # bcast A into B
+    size(B, 1) == 1 && (return size(A)) # bcast B into A
+    throw(DimensionMismatch("Got mismatched dimensions $(size(A)), $(size(B))"))
+end
+function _combinesizes(A::Transpose{<:Any, <:AbstractGBVector}, B::Transpose{<:Any, <:AbstractGBVector})
+    size(A) == size(B) && (return size(A)) # same size 
+    size(A, 2) == 1 && (return size(B)) # bcast A into B
+    size(B, 2) == 1 && (return size(A)) # bcast B into A
+    throw(DimensionMismatch("Got mismatched dimensions $(size(A)), $(size(B))"))
 end
+# Outer products (dot is done by mul[!])
+function _combinesizes(A::AbstractGBVector, B::Transpose{<:Any, <:AbstractGBVector})
+    return (size(A, 1), size(B, 2))
+end
+function _combinesizes(A::Transpose{<:Any, <:AbstractGBVector}, B::AbstractGBVector)
+    return (size(B, 1), size(A, 2))
+end
+
+function _combinesizes(A::GBMatrixOrTranspose, v::AbstractGBVector)
+    length(v)
+    size(A, 1) == size(v, 1) && (return size(A))
+    throw(DimensionMismatch("Got mismatched dimensions $(size(A, 1)) and $(size(v, 1))"))
+end
+function _combinesizes(v::AbstractGBVector, A::GBMatrixOrTranspose)
+    size(A, 1) == size(v, 1) && (return size(A))
+    throw(DimensionMismatch("Got mismatched dimensions $(size(v, 1)) and $(size(A, 1))"))
+end
+
+function _combinesizes(A::GBMatrixOrTranspose, v::Transpose{<:Any, <:AbstractGBVector})
+    size(A, 2) == size(v, 2) && (return size(A))
+    throw(DimensionMismatch("Got mismatched dimensions $(size(A, 2)) and $(size(v, 2))"))
+end
+function _combinesizes(v::Transpose{<:Any, <:AbstractGBVector}, A::GBMatrixOrTranspose)
+    size(A, 2) == size(v, 2) && (return size(A))
+    throw(DimensionMismatch("Got mismatched dimensions $(size(v, 2)) and $(size(A, 2))"))
+end
+
+_combinesizes(A::GBMatrixOrTranspose, B::GBMatrixOrTranspose) = size(A) == size(B) ? size(A) : 
+    throw(DimensionMismatch("Got mismatched dimensions: $(size(A)) and $(size(B))"))

src/operations/broadcasts.jl

@@ -26,6 +26,8 @@ Base.BroadcastStyle(::Type{<:AbstractGBVector}) = GBVectorStyle()
 Base.BroadcastStyle(::Type{<:AbstractGBMatrix}) = GBMatrixStyle()
 Base.BroadcastStyle(::Type{<:Transpose{T, <:AbstractGBMatrix} where T}) = GBMatrixStyle()
 Base.BroadcastStyle(::Type{<:Adjoint{T, <:AbstractGBMatrix} where T}) = GBMatrixStyle()
+Base.BroadcastStyle(::Type{<:Transpose{T, <:AbstractGBVector} where T}) = GBVectorStyle()
+Base.BroadcastStyle(::Type{<:Adjoint{T, <:AbstractGBVector} where T}) = GBVectorStyle()
 
 #
 GBVectorStyle(::Val{0}) = GBVectorStyle()
@@ -91,14 +93,6 @@ modifying(::typeof(emul)) = emul!
         if right isa Broadcast.Broadcasted
             right = copy(right)
         end
-        if left isa AbstractVector && right isa GBMatrixOrTranspose && 
-            !(size(left, 1) == size(right, 1) && size(left, 2) == size(right, 2))
-            return *(Diagonal(left), right, (any, f))
-        end
-        if left isa GBMatrixOrTranspose && right isa Transpose{<:Any, <:AbstractVector} && 
-            !(size(left, 1) == size(right, 1) && size(left, 2) == size(right, 2))
-            return *(left, Diagonal(right), (any, f))
-        end
         if left isa StridedArray
             left = pack(left; fill = right isa GBArrayOrTranspose ? getfill(right) : nothing)
         end
@@ -177,14 +171,6 @@ mutatingop(::typeof(apply)) = apply!
                     # If they're further nested broadcasts we can't fuse them, so just copy.
                     subargleft isa Broadcast.Broadcasted && (subargleft = copy(subargleft))
                     subargright isa Broadcast.Broadcasted && (subargright = copy(subargright))
-                    if left isa AbstractVector && right isa GBMatrixOrTranspose && 
-                        !(size(left, 1) == size(right, 1) && size(left, 2) == size(right, 2))
-                        return *(Diagonal(left), right, (any, f); accum)
-                    end
-                    if left isa GBMatrixOrTranspose && right isa Transpose{<:Any, <:AbstractVector} && 
-                        !(size(left, 1) == size(right, 1) && size(left, 2) == size(right, 2))
-                        return *(left, Diagonal(right), (any, f); accum)
-                    end
                     if subargleft isa StridedArray
                         subargleft = pack(subargleft; fill = subargright isa GBArrayOrTranspose ? getfill(right) : 0)
                     end
@@ -365,4 +351,69 @@ function Base.materialize!(
     return setindex!(A, bc.args[begin], :)
 end 
 
-Base.Broadcast.broadcasted(::Type{T}, A::AbstractGBArray) where T = LinearAlgebra.copy_oftype(A, T)
+Base.Broadcast.broadcasted(::Type{T}, A::AbstractGBArray) where T = LinearAlgebra.copy_oftype(A, T)
+
+# This is overly verbose, perhaps a macro?
+# return an operator that swaps the order of the operands.
+# * -> *, first -> second, second -> first, - -> rminus, etc.
+_swapop(op) = throw(ArgumentError("Cannot swap order of operands automatically. Swap the order of the broadcast statement or overload `_swapop`"))
+_swapop(::typeof(first)) = second
+_swapop(::typeof(second)) = first
+
+_swapop(::typeof(any)) = any
+
+_swapop(::typeof(pair)) = pair
+
+_swapop(::typeof(+)) = +
+_swapop(::typeof(-)) = rminus
+_swapop(::typeof(rminus)) = -
+
+_swapop(::typeof(*)) = *
+_swapop(::typeof(/)) = \
+_swapop(::typeof(\)) = /
+
+# ^ / POW doesn't have an equivalent builtin... Error for now.
+
+_swapop(::typeof(iseq)) = iseq
+_swapop(::typeof(isne)) = isne
+
+_swapop(::typeof(min)) = min
+_swapop(::typeof(max)) = max
+
+_swapop(::typeof(isgt)) = isle
+_swapop(::typeof(isle)) = isgt
+
+_swapop(::typeof(isge)) = islt
+_swapop(::typeof(islt)) = isge
+
+_swapop(::typeof(∨)) = ∨
+_swapop(::typeof(∧)) = ∧
+
+_swapop(::typeof(lxor)) = lxor
+_swapop(::typeof(xnor)) = xnor
+
+_swapop(::typeof(==)) = ==
+_swapop(::typeof(!=)) = !=
+
+_swapop(::typeof(>)) = <=
+_swapop(::typeof(<=)) = >
+_swapop(::typeof(<)) = >=
+_swapop(::typeof(>=)) = <
+
+# I'm not going to bother with the trig/mod/sign/complex/etc. If you need them please open an issue.
+
+_swapop(::typeof(|)) = |
+_swapop(::typeof(&)) = &
+_swapop(::typeof(⊻)) = ⊻
+_swapop(::typeof(bxnor)) = bxnor
+# bshift has no obvious equivalent in the builtins
+
+_swapop(::typeof(firsti0)) = secondi0
+_swapop(::typeof(secondi0)) = firsti0
+_swapop(::typeof(firsti)) = secondi
+_swapop(::typeof(secondi)) = firsti
+
+_swapop(::typeof(firstj0)) = secondj0
+_swapop(::typeof(secondj0)) = firstj0
+_swapop(::typeof(firstj)) = secondj
+_swapop(::typeof(secondj)) = firstj