AxB.py

from dataclasses import dataclass
from numpy import array
import typing

class GrB_NULL:
    pass

@dataclass
class GrB_Descriptor:
    pass

@dataclass
class GrB_BinaryOp:
    op: typing.Callable

    def __call__(self, first, second):
        return self.op(first, second)
        
@dataclass
class GrB_Monoid:
    op: GrB_BinaryOp
    identity: object

@dataclass
class GrB_Semiring:
    add: GrB_Monoid
    multiply: GrB_BinaryOp

@dataclass
class GrB_Matrix:
    store: array

@dataclass
class GrB_Vector(GrB_Matrix):

    def __post_init__(self):
        assert len(self.store.shape) == 1, "Vector storage must be 1-dimensional."

def GrB_vxm(
        u: GrB_Vector,
        A: GrB_Matrix,
        w: GrB_Vector = GrB_NULL,
        mask: GrB_Vector = GrB_NULL,
        accum: GrB_BinaryOp = GrB_NULL,
        semiring: GrB_Semiring = GrB_NULL,
        desc: GrB_Descriptor = GrB_NULL
) -> GrB_Vector:
    """ w'<Mask> = accum (w, u'*A) """
    print("{}'<{}> = {} ({}, {}'*{})".format(w, mask, accum, w, u, A))
    return GrB_Vector(u.store @ A.store)

a = GrB_Vector(array([1,2,3]))
B = GrB_Matrix(array([[1,2,3], [4,5,6], [7,8,9]]))

C = GrB_vxm(a, B)
print(C)

# Output:
#
# <class 'pygraphblas.AxB.GrB_NULL'>'<<class 'pygraphblas.AxB.GrB_NULL'>> = <class 'pygraphblas.AxB.GrB_NULL'> (<class 'pygraphblas.AxB.GrB_NULL'>, GrB_Vector(store=array([1, 2, 3]))'*GrB_Matrix(store=array([[1, 2, 3],
#        [4, 5, 6],
#        [7, 8, 9]])))
# GrB_Vector(store=array([30, 36, 42]))