add idxunaryops back finally

Author: rayegun

src/SuiteSparseGraphBLAS.jl

@@ -38,6 +38,7 @@ include("operators/libgbops.jl")
 
 include("gbtypes.jl")
 include("types.jl")
+include("scalar.jl")
 include("mem.jl")
 
 
@@ -55,11 +56,11 @@ using .UnaryOps
 using .BinaryOps
 using .Monoids
 using .Semirings
+using .IndexUnaryOps
 
 include("indexutils.jl")
 # 
 include("operations/extract.jl")
-include("scalar.jl")
 include("gbvector.jl")
 include("gbmatrix.jl")
 include("abstractgbarray.jl")
@@ -159,8 +160,6 @@ function __init__()
     for type ∈ valid_vec
         Base.unsafe_convert(LibGraphBLAS.GrB_Type, gbtype(type))
     end
-    # Eagerly load selectops constants.
-    _loadselectops()
     ALL.p = load_global("GrB_ALL", LibGraphBLAS.GrB_Index)
     # Set printing done by SuiteSparse:GraphBLAS to base-1 rather than base-0.
     gbset(BASE1, 1)

src/chainrules/selectrules.jl

@@ -2,7 +2,7 @@
 function frule(
     (_, _, ΔA)::Tuple,
     ::typeof(select),
-    op::Union{Function, SelectUnion},
+    op::Function,
     A::AbstractGBArray
 )
     Ω = select(op, A)
@@ -14,7 +14,7 @@ end
 function frule(
     (_, _, ΔA, _)::Tuple,
     ::typeof(select),
-    op::Union{Function, SelectUnion},
+    op::Function,
     A::AbstractGBArray,
     thunk::Union{GBScalar, Nothing, valid_union}
 )
@@ -25,7 +25,7 @@ end
 
 function rrule(
     ::typeof(select),
-    op::Union{Function, SelectUnion},
+    op::Function,
     A::AbstractGBArray
 )
     out = select(op, A)
@@ -38,7 +38,7 @@ end
 
 function rrule(
     ::typeof(select),
-    op::Union{Function, SelectUnion},
+    op::Function,
     A::AbstractGBArray,
     thunk::Union{GBScalar, Nothing, valid_union}
 )

src/operations/select.jl

@@ -1,43 +1,56 @@
 # TODO: update to modern op system.
 
+function defaultselectthunk(op, T)
+    if op ∈ (rowindex, colindex, diagindex)
+        return one(Int64)
+    elseif op ∈ (tril, triu, diag, offdiag)
+        return zero(Int64)
+    elseif op === ==
+        return zero(T)
+    else
+        throw(ArgumentError("You must pass `thunk` to select for this function."))
+    end
+end
+
 "In place version of `select`."
 function select!(
     op,
     C::GBVecOrMat,
-    A::GBArrayOrTranspose,
-    thunk = nothing;
+    A::GBArrayOrTranspose{T},
+    thunk::TH = defaultselectthunk(op, T);
     mask = nothing,
     accum = nothing,
     desc = nothing
-)
+) where {T, TH}
+    op ∈ (rowindex, colindex, diagindex, tril, triu, diag, offdiag) && 
+        (thunk = convert(Int64, thunk))
     _canbeoutput(C) || throw(ShallowException())
-    op = SelectOp(op)
+    op = indexunaryop(op, T, TH)
     desc = _handledescriptor(desc; out=C, in1=A)
     mask = _handlemask!(desc, mask)
-    thunk === nothing && (thunk = C_NULL)
     accum = _handleaccum(accum, storedeltype(C))
-    if thunk isa Number
-        thunk = GBScalar(thunk)
-    end
-    @wraperror LibGraphBLAS.GxB_Matrix_select(C, mask, accum, op, parent(A), thunk, desc)
+    @wraperror LibGraphBLAS.GrB_Matrix_select_Scalar(C, mask, accum, op, parent(A), GBScalar(thunk), desc)
     return C
 end
 
-function select!(op, A::GBArrayOrTranspose, thunk = nothing; mask = nothing, accum = nothing, desc = nothing)
+function select!(
+    op, A::GBArrayOrTranspose{T}, thunk = defaultselectthunk(op, T); 
+    mask = nothing, accum = nothing, desc = nothing
+) where T
     return select!(op, A, A, thunk; mask, accum, desc)
 end
 
 """
-    select(op::Union{Function, SelectUnion}, A::GBArrayOrTranspose; kwargs...)::GBArrayOrTranspose
-    select(op::Union{Function, SelectUnion}, A::GBArrayOrTranspose, thunk; kwargs...)::GBArrayOrTranspose
+    select(op::Function, A::GBArrayOrTranspose; kwargs...)::GBArrayOrTranspose
+    select(op::Function, A::GBArrayOrTranspose, thunk; kwargs...)::GBArrayOrTranspose
 
 Return a `GBArray` whose elements satisfy the predicate defined by `op`.
 Some SelectOps or functions may require an additional argument `thunk`, for use in
     comparison operations such as `C[i,j] = A[i,j] >= thunk ? A[i,j] : nothing`, which is
     performed by `select(>, A, thunk)`.
 
 # Arguments
-- `op::Union{Function, SelectUnion}`: A select operator from the SelectOps submodule.
+- `op::Function`: A select operator from the SelectOps submodule.
 - `A::GBArrayOrTranspose`
 - `thunk::Union{GBScalar, nothing, valid_union}`: Optional value used to evaluate `op`.
 
@@ -53,19 +66,19 @@ Some SelectOps or functions may require an additional argument `thunk`, for use
 """
 function select(
     op,
-    A::GBArrayOrTranspose,
-    thunk = nothing;
+    A::GBArrayOrTranspose{T},
+    thunk::TH = defaultselectthunk(op, T);
     mask = nothing,
     accum = nothing,
     desc = nothing
-)
-    op = SelectOp(op)
-    C = similar(A)
+) where {T, TH}
+    op = indexunaryop(op, T, TH)
+    C = similar(A) # we keep the same type!! not the ztype of op.
     select!(op, C, A, thunk; accum, mask, desc)
     return C
 end
 
 LinearAlgebra.tril(A::GBArrayOrTranspose, k::Integer = 0) = select(tril, A, k)
 LinearAlgebra.triu(A::GBArrayOrTranspose, k::Integer = 0) = select(triu, A, k)
-SparseArrays.dropzeros(A::GBArrayOrTranspose) = select(nonzeros, A)
-SparseArrays.dropzeros!(A::GBArrayOrTranspose) = select!(nonzeros, A)
+SparseArrays.dropzeros(A::GBArrayOrTranspose{T}) where T = select(!=, A, zero(T))
+SparseArrays.dropzeros!(A::GBArrayOrTranspose{T}) where T = select!(!=, A, zero(T))

src/operators/binaryops.jl

@@ -25,8 +25,8 @@ end
 SuiteSparseGraphBLAS.binaryop(f::F, ::Type{X}, ::Type{X}) where {F, X} = fallback_binaryop(f, X, X)
 
 function SuiteSparseGraphBLAS.binaryop(
-    f::F, ::Type{X}, ::Type{Y}
-) where {F, X, Y}
+    f, ::Type{X}, ::Type{Y}
+) where {X, Y}
     P = promote_type(X, Y)
     if isconcretetype(P) && (X <: valid_union && Y <: valid_union)
         return binaryop(f, P, P)
@@ -41,8 +41,8 @@ SuiteSparseGraphBLAS.binaryop(f, ::GBArrayOrTranspose{T}, ::Type{U}) where {T, U
 SuiteSparseGraphBLAS.binaryop(f, ::Type{T}, ::GBArrayOrTranspose{U}) where {T, U} = binaryop(f, T, U)
 
 SuiteSparseGraphBLAS.binaryop(f, type) = binaryop(f, type, type)
-SuiteSparseGraphBLAS.binaryop(op::TypedBinaryOperator, x...) = op
-
+SuiteSparseGraphBLAS.binaryop(op::TypedBinaryOperator, ::Type{X}, ::Type{Y}) where {X, Y} = op
+SuiteSparseGraphBLAS.binaryop(op::TypedBinaryOperator, ::Type{X}) where X = op
 SuiteSparseGraphBLAS.juliaop(op::TypedBinaryOperator) = op.fn
 
 # TODO, clean up this function, it allocates typedop and is otherwise perhaps a little slow.