Ctrl+K

seg_matmul

Contents

seg_matmul#

class brainpy.math.sparse.seg_matmul(A, B)[source]#

Sparse matrix multiplication.

\[y = A @ B\]

where \(A\) or \(B\) is a sparse matrix. \(A\) and \(B\) cannot be both sparse.

Examples

>>> import brainpy.math as bm

  1. when the left matrix \(A\) is a sparse matrix with the shape of \((N, M)\),

>>> # A is a sparse matrix (3, 4):
>>> #   [[0, 2, 0, 4],
>>> #    [1, 0, 0, 0],
>>> #    [0, 3, 0, 2]]
>>> values = bm.asarray([2, 4, 1, 3, 2])
>>> rows = bm.asarray([0, 0, 1, 2, 2])
>>> cols = bm.asarray([1, 3, 0, 1, 3])
>>> sparse = {'data': values, 'index': (rows, cols), 'shape': (3, 4)}
>>> B = bm.arange(4)
>>> bm.sparse.sparse_matmul(sparse, B)
ArrayType([14,  0,  9], dtype=int32)
>>> B = bm.random.rand(4, 3)
>>> bm.sparse.sparse_matmul(sparse, B)
ArrayType([[3.8331761 , 1.3708692 , 4.510223  ],
          [0.9960836 , 0.37550318, 0.7370341 ],
          [2.3700516 , 0.7574289 , 4.1124535 ]], dtype=float32)

  1. when the right matrix \(B\) is a sparse matrix with the shape of \((M, K)\),

>>> A = bm.arange(3)
>>> bm.sparse.sparse_matmul(A, sparse)
ArrayType([1, 6, 0, 4], dtype=int32)
>>> A = bm.random.rand(2, 3)
>>> bm.sparse.sparse_matmul(A, sparse)
ArrayType([[0.438388  , 1.4346815 , 0.        , 2.361964  ],
          [0.9171978 , 1.1214957 , 0.        , 0.90534496]],  dtype=float32)
Parameters:

  • A (tensor, sequence) – The dense or sparse matrix with the shape of \((N, M)\).

  • B (tensor, sequence) – The dense or sparse matrix with the shape of \((M, K)\).

Returns:

results – The tensor with the shape of \((N, K)\).

Return type:

ArrayType

Contents