with function/ctxvarx, infer output type better

Author: Wimmerer

Project.toml

@@ -8,6 +8,7 @@ CEnum = "fa961155-64e5-5f13-b03f-caf6b980ea82"
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
+ContextVariablesX = "6add18c4-b38d-439d-96f6-d6bc489c04c5"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"

src/SuiteSparseGraphBLAS.jl

@@ -7,7 +7,7 @@ using MacroTools
 using LinearAlgebra
 using Random: randsubseq, default_rng, AbstractRNG, GLOBAL_RNG
 using CEnum
-
+using ContextVariablesX
 include("abstracts.jl")
 include("libutils.jl")
 include("lib/LibGraphBLAS.jl")
@@ -57,6 +57,18 @@ const MonoidBinaryOrRig = Union{
     AbstractBinaryOp,
     AbstractMonoid
 }
+
+const OperatorUnion = Union{
+    AbstractOp,
+    GrBOp
+}
+
+@contextvar ctxop::OperatorUnion
+@contextvar ctxmask::Union{GBArray, Ptr} = C_NULL
+@contextvar ctxaccum::Union{BinaryUnion, Ptr} = C_NULL
+@contextvar ctxdesc::Descriptor
+include("with.jl")
+
 include("scalar.jl")
 include("vector.jl")
 include("matrix.jl")
@@ -86,12 +98,13 @@ export GBScalar, GBVector, GBMatrix #arrays
 export clear!, extract, extract!, subassign!, assign! #array functions
 
 #operations
-export mul, select, select!, eadd, eadd!, emul, emul!, apply, apply!, gbtranspose, gbtranspose!
-export multiply
+export mul, select, select!, eadd, eadd!, emul, emul!, map, map!, gbtranspose, gbtranspose!
 # Reexports.
 export diag, Diagonal, mul!, kron, kron!, transpose
 export nnz, sprand, findnz
 
+#with function
+export with
 
 function __init__()
     _load_globaltypes()
@@ -107,4 +120,8 @@ function __init__()
         libgb.GrB_finalize()
     end
 end
+
+#We need to do this after __init__
+
+
 end #end of module

src/operations/ewise.jl

@@ -19,7 +19,7 @@ union equivalent see [`eadd!`](@ref).
 - `mask::Union{Ptr{Nothing}, GBMatrix} = C_NULL`: optional mask.
 - `accum::Union{Ptr{Nothing}, AbstractBinaryOp} = C_NULL`: binary accumulator operation
     where `C[i,j] = accum(C[i,j], T[i,j])` where T is the result of this function before accum is applied.
-- `desc::Descriptor = Descriptors.NULL`
+- `desc = nothing`
 """
 function emul! end
 
@@ -43,7 +43,7 @@ union equivalent see [`eadd`](@ref).
 - `mask::Union{Ptr{Nothing}, GBMatrix} = C_NULL`: optional mask.
 - `accum::Union{Ptr{Nothing}, AbstractBinaryOp} = C_NULL`: binary accumulator operation
     where `C[i,j] = accum(C[i,j], T[i,j])` where T is the result of this function before accum is applied.
-- `desc::Descriptor = Descriptors.NULL`
+- `desc = nothing`
 
 # Returns
 - `GBArray`: Output `GBVector` or `GBMatrix` whose eltype is determined by the `eltype` of
@@ -55,11 +55,13 @@ function emul!(
     w::GBVector,
     u::GBVector,
     v::GBVector,
-    op::MonoidBinaryOrRig = BinaryOps.TIMES;
-    mask = C_NULL,
-    accum = C_NULL,
-    desc::Descriptor = Descriptors.NULL
+    op = nothing;
+    mask = nothing,
+    accum = nothing,
+    desc = nothing
 )
+    op, mask, accum, desc = _handlectx(op, mask, accum, desc, BinaryOps.TIMES)
+
     size(w) == size(u) == size(v) || throw(DimensionMismatch())
     op = getoperator(op, optype(u, v))
     accum = getoperator(accum, eltype(w))
@@ -72,17 +74,20 @@ function emul!(
     elseif op isa libgb.GrB_BinaryOp
         libgb.GrB_Vector_eWiseMult_BinaryOp(w, mask, accum, op, u, v, desc)
         return w
+    else
+        throw(ArgumentError("$op is not a valid monoid binary op or semiring."))
     end
 end
 
 function emul(
     u::GBVector,
     v::GBVector,
-    op::MonoidBinaryOrRig = BinaryOps.TIMES;
-    mask = C_NULL,
-    accum = C_NULL,
-    desc::Descriptor = Descriptors.NULL
+    op = nothing;
+    mask = nothing,
+    accum = nothing,
+    desc = nothing
 )
+    op = _handlectx(op, ctxop, BinaryOps.TIMES)
     if op isa GrBOp
         t = ztype(op)
     else
@@ -96,11 +101,12 @@ function emul!(
     C::GBMatrix,
     A::GBMatOrTranspose,
     B::GBMatOrTranspose,
-    op::MonoidBinaryOrRig = BinaryOps.TIMES;
-    mask = C_NULL,
-    accum = C_NULL,
-    desc::Descriptor = Descriptors.NULL
+    op = nothing;
+    mask = nothing,
+    accum = nothing,
+    desc = nothing
 )
+    op, mask, accum, desc = _handlectx(op, mask, accum, desc, BinaryOps.TIMES)
     size(C) == size(A) == size(B) || throw(DimensionMismatch())
     A, desc, B = _handletranspose(A, desc, B)
     op = getoperator(op, optype(A, B))
@@ -114,17 +120,20 @@ function emul!(
     elseif op isa libgb.GrB_BinaryOp
         libgb.GrB_Matrix_eWiseMult_BinaryOp(C, mask, accum, op, A, B, desc)
         return C
+    else
+        throw(ArgumentError("$op is not a valid monoid binary op or semiring."))
     end
 end
 
 function emul(
     A::GBMatOrTranspose,
     B::GBMatOrTranspose,
-    op::MonoidBinaryOrRig = BinaryOps.TIMES;
-    mask = C_NULL,
-    accum = C_NULL,
-    desc::Descriptor = Descriptors.NULL
+    op = nothing;
+    mask = nothing,
+    accum = nothing,
+    desc = nothing
 )
+    op = _handlectx(op, ctxop, BinaryOps.TIMES)
     if op isa GrBOp
         t = ztype(op)
     else
@@ -155,7 +164,7 @@ intersection equivalent see [`eadd!`](@ref).
 - `mask::Union{Ptr{Nothing}, GBMatrix} = C_NULL`: optional mask.
 - `accum::Union{Ptr{Nothing}, AbstractBinaryOp} = C_NULL`: binary accumulator operation
     where `C[i,j] = accum(C[i,j], T[i,j])` where T is the result of this function before accum is applied.
-- `desc::Descriptor = Descriptors.NULL`
+- `desc = nothing`
 """
 function eadd! end
 
@@ -179,7 +188,7 @@ intersection equivalent see [`emul`](@ref).
 - `mask::Union{Ptr{Nothing}, GBMatrix} = C_NULL`: optional mask.
 - `accum::Union{Ptr{Nothing}, AbstractBinaryOp} = C_NULL`: binary accumulator operation
     where `C[i,j] = accum(C[i,j], T[i,j])` where T is the result of this function before accum is applied.
-- `desc::Descriptor = Descriptors.NULL`
+- `desc = nothing`
 
 # Returns
 - `GBArray`: Output `GBVector` or `GBMatrix` whose eltype is determined by the `eltype` of
@@ -191,11 +200,12 @@ function eadd!(
     w::GBVector,
     u::GBVector,
     v::GBVector,
-    op::MonoidBinaryOrRig = BinaryOps.PLUS;
-    mask = C_NULL,
-    accum = C_NULL,
-    desc::Descriptor = Descriptors.NULL
+    op = nothing;
+    mask = nothing,
+    accum = nothing,
+    desc = nothing
 )
+    op, mask, accum, desc = _handlectx(op, mask, accum, desc, BinaryOps.PLUS)
     size(w) == size(u) == size(v) || throw(DimensionMismatch())
     op = getoperator(op, optype(u, v))
     accum = getoperator(accum, eltype(w))
@@ -208,17 +218,20 @@ function eadd!(
     elseif op isa libgb.GrB_BinaryOp
         libgb.GrB_Vector_eWiseAdd_BinaryOp(w, mask, accum, op, u, v, desc)
         return w
+    else
+        throw(ArgumentError("$op is not a valid monoid binary op or semiring."))
     end
 end
 
 function eadd(
     u::GBVector,
     v::GBVector,
-    op::MonoidBinaryOrRig = BinaryOps.PLUS;
-    mask = C_NULL,
-    accum = C_NULL,
-    desc::Descriptor = Descriptors.NULL
+    op = nothing;
+    mask = nothing,
+    accum = nothing,
+    desc = nothing
 )
+    op, mask, accum, desc = _handlectx(op, mask, accum, desc, BinaryOps.PLUS)
     if op isa GrBOp
         t = ztype(op)
     else
@@ -232,11 +245,12 @@ function eadd!(
     C::GBMatrix,
     A::GBMatOrTranspose,
     B::GBMatOrTranspose,
-    op::MonoidBinaryOrRig = BinaryOps.PLUS;
-    mask = C_NULL,
-    accum = C_NULL,
-    desc::Descriptor = Descriptors.NULL
+    op = nothing;
+    mask = nothing,
+    accum = nothing,
+    desc = nothing
 )
+    op, mask, accum, desc = _handlectx(op, mask, accum, desc, BinaryOps.PLUS)
     size(C) == size(A) == size(B) || throw(DimensionMismatch())
     A, desc, B = _handletranspose(A, desc, B)
     op = getoperator(op, optype(A, B))
@@ -251,18 +265,19 @@ function eadd!(
         libgb.GrB_Matrix_eWiseAdd_BinaryOp(C, mask, accum, op, A, B, desc)
         return C
     else
-        error("Unreachable")
+        throw(ArgumentError("$op is not a valid monoid binary op or semiring."))
     end
 end
 
 function eadd(
     A::GBMatOrTranspose,
     B::GBMatOrTranspose,
-    op::MonoidBinaryOrRig = BinaryOps.PLUS;
-    mask = C_NULL,
-    accum = C_NULL,
-    desc::Descriptor = Descriptors.NULL
+    op = nothing;
+    mask = nothing,
+    accum = nothing,
+    desc = nothing
 )
+    op, mask, accum, desc = _handlectx(op, mask, accum, desc, BinaryOps.PLUS)
     if op isa GrBOp
         t = ztype(op)
     else
@@ -271,57 +286,3 @@ function eadd(
     C = GBMatrix{t}(size(A))
     return eadd!(C, A, B, op; mask, accum, desc)
 end
-
-# Note well: `.*` and `.+` have clear counterparts in the language of GraphBLAS:
-# edgewiseAdd and edgewiseMul. These do not necessarily have the same semantics though.
-# edgewiseAdd and edgewiseMul might better be described as edgewiseUnion and
-# edgewiseIntersection respectively, and then `op` is applied at materialized indices.
-#
-# So the plan is thus: `.*` and `.+` will have the Union and Intersection semantics *with*
-# the default ops of `*` and `+` respectively. *However*, they have `op` kwargs, which
-# may be used with a macro later on down the line to override the default ops.
-function Base.broadcasted(
-    ::typeof(+),
-    u::GBVector,
-    v::GBVector,
-    op::MonoidBinaryOrRig = BinaryOps.PLUS;
-    mask = C_NULL,
-    accum = C_NULL,
-    desc::Descriptor = Descriptors.NULL
-)
-    return eadd(u, v, op; mask, accum, desc)
-end
-function Base.broadcasted(
-    ::typeof(*),
-    u::GBVector,
-    v::GBVector,
-    op::MonoidBinaryOrRig = BinaryOps.TIMES;
-    mask = C_NULL,
-    accum = C_NULL,
-    desc::Descriptor = Descriptors.NULL
-)
-    return emul(u, v, op; mask, accum, desc)
-end
-
-function Base.broadcasted(
-    ::typeof(+),
-    A::GBMatOrTranspose,
-    B::GBMatOrTranspose,
-    op::MonoidBinaryOrRig = BinaryOps.PLUS;
-    mask = C_NULL,
-    accum = C_NULL,
-    desc::Descriptor = Descriptors.NULL
-)
-    return eadd(A, B, op; mask, accum, desc)
-end
-function Base.broadcasted(
-    ::typeof(*),
-    A::GBMatOrTranspose,
-    B::GBMatOrTranspose,
-    op::MonoidBinaryOrRig = BinaryOps.PLUS;
-    mask = C_NULL,
-    accum = C_NULL,
-    desc::Descriptor = Descriptors.NULL
-)
-    return emul(A, B, op; mask, accum, desc)
-end

src/operations/kronecker.jl

@@ -2,32 +2,40 @@ function LinearAlgebra.kron!(
     C::GBMatOrTranspose,
     A::GBMatOrTranspose,
     B::GBMatOrTranspose,
-    op::MonoidBinaryOrRig = BinaryOps.TIMES;
-    mask = C_NULL,
-    accum = C_NULL,
-    desc::Descriptor = Descriptors.NULL
+    op = nothing;
+    mask = nothing,
+    accum = nothing,
+    desc = nothing
 )
+    op, mask, accum, desc = _handlectx(op, mask, accum, desc, BinaryOps.TIMES)
     op = getoperator(op, optype(A, B))
     A, desc, B = _handletranspose(A, desc, B)
     accum = getoperator(accum, eltype(C))
-    if op isa Union{AbstractBinaryOp, libgb.GrB_BinaryOp}
+    if op isa libgb.GrB_BinaryOp
         libgb.GxB_kron(C, mask, accum, op, A, B, desc)
-    elseif op isa Union{AbstractMonoid, libgb.GrB_Monoid}
+    elseif op isa libgb.GrB_Monoid
         libgb.GrB_Matrix_kronecker_Monoid(C, mask, accum, op, A, B, desc)
-    elseif op isa Union{AbstractSemiring, libgb.GrB_Semiring}
+    elseif op isa libgb.GrB_Semiring
         libgb.GrB_Matrix_kronecker_Semiring(C, mask, accum, op, A, B, desc)
+    else
+        throw(ArgumentError("$op is not a valid monoid binary op or semiring."))
     end
 end
 
 function LinearAlgebra.kron(
     A::GBMatOrTranspose,
     B::GBMatOrTranspose,
-    op::MonoidBinaryOrRig = BinaryOps.TIMES;
-    mask = C_NULL,
-    accum = C_NULL,
-    desc::Descriptor = Descriptors.NULL
+    op = nothing;
+    mask = nothing,
+    accum = nothing,
+    desc = nothing
 )
-    t = optype(A, B)
+    op = _handlectx(op, ctxop, BinaryOps.TIMES)
+    if op isa GrBOp
+        t = ztype(op)
+    else
+        t = optype(A, B)
+    end
     C = GBMatrix{t}(size(A,1) * size(B, 1), size(A, 2) * size(B, 2))
     kron!(C, A, B, op; mask, accum, desc)
     return C