Sparse helpers cover both sorted-index intersections and weighted sparse dot products.
import numpy as np
import numkong as nk
idx_a, idx_b = np.array([1, 3, 5, 7], dtype=np.uint32), np.array([3, 4, 5, 8], dtype=np.uint32)
intersection_size = nk.intersect(idx_a, idx_b) # len(np.intersect1d(idx_a, idx_b))
assert intersection_size == 2, "indices 3 and 5"
val_a, val_b = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32), np.array([5.0, 6.0, 7.0, 8.0], dtype=np.float32)
sparse_dot = nk.sparse_dot(idx_a, val_a, idx_b, val_b)
assert sparse_dot > 0, "weighted dot over shared indices"intersectCount the elements two sorted integer arrays share.def intersect(
a: _BufferType,
b: _BufferType,
/,
*,
capabilities: Capability | int | None = ...,
stream: int | None = ...,
) -> intSimilar to numpy.intersect1d, returning the size of the intersection rather than its elements.
capabilitiesCapability | int | None | Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them. | |
|---|---|---|
streamint | None | A GPU stream pointer, left None on the CPU. |
sparse_dotCompute the weighted dot product of two sparse vectors over their shared indices.def sparse_dot(
a_indices: _BufferType,
a_values: _BufferType,
b_indices: _BufferType,
b_values: _BufferType,
/,
*,
capabilities: Capability | int | None = ...,
stream: int | None = ...,
) -> floata_indices_BufferType | First sorted index array, uint16 or uint32. | |
|---|---|---|
a_values_BufferType | Weights of a_indices, bfloat16 or float32. | |
b_indices_BufferType | Second sorted index array, of the dtype of a_indices. | |
b_values_BufferType | Weights of b_indices, of the dtype of a_values. | |
capabilitiesCapability | int | None | Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them. | |
streamint | None | A GPU stream pointer, left None on the CPU. |