[DOC, TESTS] Add doc, tests

Author: ricardosp4

doc/reference/array_operations.rst

@@ -6,3 +6,4 @@ Operations with arrays
 
     lazy_functions
     reduction_functions
+    linear_algebra

doc/reference/linear_algebra.rst

@@ -0,0 +1,14 @@
+.. _linear_algebra:
+
+Linear Algebra
+--------------
+
+The next functions can be used for computing linear algebra operations with :ref:`NDArray <NDArray>`.
+
+.. currentmodule:: blosc2
+
+.. autosummary::
+   :toctree: autofiles/operations_with_arrays/
+   :nosignatures:
+
+    matmul

src/blosc2/__init__.py

@@ -235,6 +235,7 @@ class Tuner(Enum):
     ones,
     full,
     save,
+    matmul,
 )
 
 from .c2array import c2context, C2Array, URLPath

src/blosc2/ndarray.py

@@ -3645,23 +3645,70 @@ def sort(array: NDArray, order: str | list[str] | None = None, **kwargs: Any) ->
 
 def matmul(x1: NDArray, x2: NDArray, **kwargs: Any) -> NDArray:
     """
-    Perform matrix multiplication between two Blosc2 NDArrays.
+    Computes the matrix product between two Blosc2 NDArrays.
 
     Parameters
     ----------
     x1: `NDArray`
-        First input array.
+        The first input array.
     x2: `NDArray`
-        Second input array.
+        The second input array.
     kwargs: Any, optional
         Keyword arguments that are supported by the :func:`empty` constructor.
 
     Returns
     -------
     out: :ref:`NDArray`
-        The result matrix multiplication.
+        The matrix product of the inputs. This is a scalar only when both x1,
+        x2 are 1-d vectors.
+
+    Raises
+    ------
+    ValueError
+        If the last dimension of ``x1`` is not the same size as
+        the second-to-last dimension of ``x2``.
+
+        If a scalar value is passed in.
+
+    References
+    ----------
+    `numpy.matmul <https://numpy.org/doc/stable/reference/generated/numpy.matmul.html>`_
+
+    Examples
+    --------
+    For 2-D arrays it is the matrix product:
+
+    >>> import numpy as np
+    >>> import blosc2
+    >>> a = np.array([[1, 2],
+    ...               [3, 4]])
+    >>> nd_a = blosc2.asarray(a)
+    >>> b = np.array([[2, 3],
+    ...               [2, 1]])
+    >>> nd_b = blosc2.asarray(b)
+    >>> blosc2.matmul(nd_a, nd_b)
+    array([[ 6,  5],
+           [14, 13]])
+
+
+    For 2-D mixed with 1-D, the result is the usual.
+
+    >>> a = np.array([[1, 3],
+    ...               [0, 1]])
+    >>> nd_a = blosc2.asarray(a)
+    >>> v = np.array([1, 2])
+    >>> nd_v = blosc2.asarray(v)
+    >>> blosc2.matmul(nd_a, nd_v)
+    array([7, 2])
+    >>> blosc2.matmul(nd_v, nd_a)
+    array([1, 5])
+
     """
 
+    # Validate arguments are not scalars
+    if np.isscalar(x1) or np.isscalar(x2):
+        raise ValueError("Arguments can't be scalars.")
+
     # Promote 1D arrays to 2D if necessary
     x1_is_vector = False
     x2_is_vector = False
@@ -3673,13 +3720,14 @@ def matmul(x1: NDArray, x2: NDArray, **kwargs: Any) -> NDArray:
         x2_is_vector = True
 
     # Validate matrix multiplication compatibility
-    if x1.shape[-1] != x2.shape[0]:
+    if x1.shape[-1] != x2.shape[-2]:
         raise ValueError("Shapes are not aligned for matrix multiplication.")
 
-    n, l = x1.shape
-    _, m = x2.shape
-    p1, q1 = x1.chunks
-    p2, q2 = x2.chunks
+    n, l = x1.shape[-2:]
+    m = x2.shape[-1]
+
+    p1, q1 = x1.chunks[-2:]
+    q2 = x2.chunks[-1]
 
     result = np.zeros((n, m), dtype=x1.dtype)
     # result = blosc2.zeros(n) # TODO: file a ticket for blosc2.zeros()
@@ -3696,7 +3744,7 @@ def matmul(x1: NDArray, x2: NDArray, **kwargs: Any) -> NDArray:
                 bres[:] += np.matmul(bx1, bx2)
 
     if x1_is_vector and x2_is_vector:
-        result = result.reshape((n, m))
+        result = result[0][0]
     elif x1_is_vector:
         result = result.reshape((m,))
     elif x2_is_vector:

tests/ndarray/test_matmul.py

@@ -1,14 +1,13 @@
 import pytest
 import numpy as np
 import blosc2
-from blosc2.ndarray import matmul
 
 
 @pytest.mark.parametrize(
     ("ashape", "achunks", "ablocks"),
     [
-        ((12, 10), (6, 5), (3, 3)),
-        ((10, ), (9, ), (7, )),
+        ((12, 10), (7, 5), (3, 3)),
+        ((10,), (9,), (7,)),
     ],
 )
 @pytest.mark.parametrize(
@@ -25,18 +24,19 @@
 def test_matmul(ashape, achunks, ablocks, bshape, bchunks, bblocks, dtype):
     a = blosc2.linspace(0, 10, dtype=dtype, shape=ashape, chunks=achunks, blocks=ablocks)
     b = blosc2.linspace(0, 10, dtype=dtype, shape=bshape, chunks=bchunks, blocks=bblocks)
+    blosc2_res = blosc2.matmul(a, b)
 
     na = a[:]
     nb = b[:]
-    blosc2_res = matmul(a, b)
     np_res = np.matmul(na, nb)
+
     np.testing.assert_allclose(blosc2_res, np_res, rtol=1e-6)
 
 
 @pytest.mark.parametrize(
     ("ashape", "achunks", "ablocks"),
     [
-        ((12, 11), (6, 5), (3, 1)),
+        ((12, 11), (7, 5), (3, 1)),
         ((0, 0), (0, 0), (0, 0)),
         ((10,), (4,), (2,)),
     ],
@@ -45,13 +45,40 @@ def test_matmul(ashape, achunks, ablocks, bshape, bchunks, bblocks, dtype):
     ("bshape", "bchunks", "bblocks"),
     [
         ((1, 5), (1, 4), (1, 3)),
-        ((4, 12), (2, 4), (1, 3)),
-        ((10,), (4,), (2,)),
+        ((4, 6), (2, 4), (1, 3)),
+        ((5,), (4,), (2,)),
     ],
 )
-def test_matmul_raises(ashape, achunks, ablocks, bshape, bchunks, bblocks):
+def test_matmul_shapes(ashape, achunks, ablocks, bshape, bchunks, bblocks):
     a = blosc2.linspace(0, 10, shape=ashape, chunks=achunks, blocks=ablocks)
     b = blosc2.linspace(0, 10, shape=bshape, chunks=bchunks, blocks=bblocks)
-    if a.shape[-1] != b.shape[0]:
-        with pytest.raises(ValueError):
-            matmul(a, b)
+
+    with pytest.raises(ValueError):
+        blosc2.matmul(a, b)
+
+    with pytest.raises(ValueError):
+        blosc2.matmul(b, a)
+
+
+@pytest.mark.parametrize("scalar", [
+    5,                      # int
+    5.3,                    # float
+    1 + 2j,                 # complex
+    np.int32(5),            # NumPy int32
+    np.int64(5),            # NumPy int64
+    np.float32(5.3),        # NumPy float32
+    np.float64(5.3),        # NumPy float64
+    np.complex64(1 + 2j),   # NumPy complex64
+    np.complex128(1 + 2j),  # NumPy complex128
+])
+def test_matmul_scalars(scalar):
+    vector = blosc2.asarray(np.array([1, 2, 3]))
+
+    with pytest.raises(ValueError):
+        blosc2.matmul(scalar, vector)
+
+    with pytest.raises(ValueError):
+        blosc2.matmul(vector, scalar)
+
+    with pytest.raises(ValueError):
+        blosc2.matmul(scalar, scalar)