From 8f15669ae2e8ebe5dcab423de642c6778f615be9 Mon Sep 17 00:00:00 2001 From: Matthew A Johnson Date: Wed, 17 Jun 2026 02:29:00 +0100 Subject: [PATCH] matrix-fma-and-fancy-indexing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A `Matrix` capability release. The dense matrix type gains NumPy-style fancy indexing, comparison masks and lexicographic comparison operators, a ``where`` selector, a single-rounding ``fma`` with vector broadcasting, and ``sqrt``. Two older method names move to their NumPy spellings: ``select`` becomes ``take`` and ``clip`` adopts ``min`` / ``max`` keyword bounds. **New Features** - **Fancy indexing** — ``m[[r0, r1]]`` / ``m[[r0, r1], :]`` gather rows and ``m[:, [c0, c1]]`` gathers columns, returning a :class:`Matrix`; the matching assignment forms scatter into rows or columns with last-write-wins duplicates and all-or-nothing validation. New :meth:`Matrix.take` and :meth:`Matrix.put` expose the same gather/scatter as methods, ``put`` with an ``accumulate=True`` mode that folds duplicate indices. - **Comparison masks** — :meth:`Matrix.less`, ``less_equal``, ``greater``, ``greater_equal``, ``equal``, and ``not_equal`` return a ``1.0`` / ``0.0`` mask matrix, accepting a same-shape matrix, a scalar (including ``bool``), a ``1x1`` matrix, a broadcasting row/column vector, or a list/tuple of numbers. Distinct from the comparison operators, which return a single bool. - **Lexicographic comparison operators** — ``<`` ``<=`` ``>`` ``>=`` ``==`` ``!=`` compare element by element in row-major order and return a single :class:`bool`. ``==`` / ``!=`` are total: a shape mismatch or an uncoercible list/tuple yields ``False`` / ``True`` rather than raising, so ``matrix in some_list`` works. A ``NaN`` never decides the comparison, so an all-``NaN`` matrix compares ``==`` equal to itself. Defining value equality makes :class:`Matrix` unhashable. - **`Matrix.where(mask, a, b)`** — a NumPy-style selector taking *a* where the mask is non-zero (``NaN`` counts as non-zero) and *b* elsewhere; *a* and *b* may each be a scalar, a same-shape matrix, or a list/tuple of numbers. - **`Matrix.fma(b, c)`** — fused multiply-add computing single-rounding ``self * b + c``; *b* and *c* may be a same-shape matrix, a ``1x1`` matrix, a scalar, or a row / column vector that broadcasts against ``self``. The contraction kernel is preserved so hardware FMA still applies. Use it as an accuracy primitive — compare results with :meth:`Matrix.allclose`, never ``==``. - **`Matrix.sqrt()`** — element-wise square root (negative inputs map to ``NaN``), with an ``in_place=True`` form. **Breaking Changes** - **`Matrix.select` renamed to `Matrix.take`.** The gather method is now spelled :meth:`Matrix.take` to match NumPy and pair with the new :meth:`Matrix.put`. Replace ``m.select(indices, axis)`` with ``m.take(indices, axis)``; the signature and semantics are otherwise unchanged. - **`Matrix.clip` bounds are now `min` / `max` keywords.** The signature changes from ``clip(min_or_maxval, maxval=None)`` to ``clip(min=None, max=None)``, matching :func:`numpy.clip`. Either bound may be omitted to leave that side unbounded: ``m.clip(min=0.0)`` clamps only below, ``m.clip(max=255.0)`` only above. **Documentation** - Expanded the :doc:`api` matrix surface for the new indexing, masking, comparison, ``where``, ``fma``, and ``sqrt`` methods via the ``__init__.pyi`` stub docstrings, including the totality, ``NaN``, and broadcasting rules. **Tests** - Extensive `test_matrix.py` additions covering fancy-index gather/scatter, ``take`` / ``put`` (including accumulate and all-or-nothing validation), the comparison masks, lexicographic operators (totality, ``NaN``, reflected-scalar, and list/tuple/bool coercion edge cases), ``where`` selection and value propagation, and ``fma`` row/column broadcasting. **Internal** - New `bench_fma` and `bench_take` micro-benchmarks in `scripts/bench_matrix.py`. The `examples/boids.py` demo migrates from ``select`` to ``take``. Signed-off-by: Matthew A Johnson Signed-off-by: Matthew A Johnson --- .flake8 | 3 +- CHANGELOG.md | 74 + CITATION.cff | 4 +- examples/boids.py | 4 +- pyproject.toml | 2 +- scripts/bench_matrix.py | 117 +- sphinx/source/api.rst | 2 +- sphinx/source/conf.py | 2 +- src/bocpy/__init__.pyi | 400 ++++- src/bocpy/_math.c | 1258 +++++++++++++-- templates/c_abi_consumer/pyproject.toml | 6 +- test/test_matrix.py | 1917 +++++++++++++++++++++-- 12 files changed, 3565 insertions(+), 224 deletions(-) diff --git a/.flake8 b/.flake8 index 32ca15a..4e143ba 100644 --- a/.flake8 +++ b/.flake8 @@ -17,7 +17,8 @@ extend-exclude = scripts/_vendored_warehouse_wheel.py per-file-ignores = test/*:D102,D103,D403 # .pyi stubs intentionally: + # A002: arguments shadowing Python builtins (Matrix.clip min/max bounds) # A003: methods shadowing Python builtins (Matrix.sum/min/max/abs/round) # D402: Sphinx-style docstrings start with the function signature # N802: NumPy-style ``Matrix.T`` transpose property is upper-case - *.pyi:A003,D402,N802 \ No newline at end of file + *.pyi:A003,D402,N802,A002 \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index 24f675c..bf0693b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,77 @@ +## 2026-06-17 - Version 0.12.0 +A `Matrix` capability release. The dense matrix type gains NumPy-style fancy +indexing, comparison masks and lexicographic comparison operators, a +``where`` selector, a single-rounding ``fma`` with vector broadcasting, and +``sqrt``. Two older method names move to their NumPy spellings: ``select`` +becomes ``take`` and ``clip`` adopts ``min`` / ``max`` keyword bounds. + +**New Features** + +- **Fancy indexing** — ``m[[r0, r1]]`` / ``m[[r0, r1], :]`` gather rows and + ``m[:, [c0, c1]]`` gathers columns, returning a :class:`Matrix`; the + matching assignment forms scatter into rows or columns with last-write-wins + duplicates and all-or-nothing validation. New :meth:`Matrix.take` and + :meth:`Matrix.put` expose the same gather/scatter as methods, ``put`` with + an ``accumulate=True`` mode that folds duplicate indices. +- **Comparison masks** — :meth:`Matrix.less`, ``less_equal``, ``greater``, + ``greater_equal``, ``equal``, and ``not_equal`` return a ``1.0`` / ``0.0`` + mask matrix, accepting a same-shape matrix, a scalar (including ``bool``), a + ``1x1`` matrix, a broadcasting row/column vector, or a list/tuple of + numbers. Distinct from the comparison operators, which return a single + bool. +- **Lexicographic comparison operators** — ``<`` ``<=`` ``>`` ``>=`` ``==`` + ``!=`` compare element by element in row-major order and return a single + :class:`bool`. ``==`` / ``!=`` are total: a shape mismatch or an + uncoercible list/tuple yields ``False`` / ``True`` rather than raising, so + ``matrix in some_list`` works. A ``NaN`` never decides the comparison, so an + all-``NaN`` matrix compares ``==`` equal to itself. Defining value equality + makes :class:`Matrix` unhashable. +- **`Matrix.where(mask, a, b)`** — a NumPy-style selector taking *a* where the + mask is non-zero (``NaN`` counts as non-zero) and *b* elsewhere; *a* and *b* + may each be a scalar, a same-shape matrix, or a list/tuple of numbers. +- **`Matrix.fma(b, c)`** — fused multiply-add computing single-rounding + ``self * b + c``; *b* and *c* may be a same-shape matrix, a ``1x1`` matrix, + a scalar, or a row / column vector that broadcasts against ``self``. The + contraction kernel is preserved so hardware FMA still applies. Use it as an + accuracy primitive — compare results with :meth:`Matrix.allclose`, never + ``==``. +- **`Matrix.sqrt()`** — element-wise square root (negative inputs map to + ``NaN``), with an ``in_place=True`` form. + +**Breaking Changes** + +- **`Matrix.select` renamed to `Matrix.take`.** The gather method is now + spelled :meth:`Matrix.take` to match NumPy and pair with the new + :meth:`Matrix.put`. Replace ``m.select(indices, axis)`` with + ``m.take(indices, axis)``; the signature and semantics are otherwise + unchanged. +- **`Matrix.clip` bounds are now `min` / `max` keywords.** The signature + changes from ``clip(min_or_maxval, maxval=None)`` to + ``clip(min=None, max=None)``, matching :func:`numpy.clip`. Either bound may + be omitted to leave that side unbounded: ``m.clip(min=0.0)`` clamps only + below, ``m.clip(max=255.0)`` only above. + +**Documentation** + +- Expanded the :doc:`api` matrix surface for the new indexing, masking, + comparison, ``where``, ``fma``, and ``sqrt`` methods via the + ``__init__.pyi`` stub docstrings, including the totality, ``NaN``, and + broadcasting rules. + +**Tests** + +- Extensive `test_matrix.py` additions covering fancy-index gather/scatter, + ``take`` / ``put`` (including accumulate and all-or-nothing validation), + the comparison masks, lexicographic operators (totality, ``NaN``, + reflected-scalar, and list/tuple/bool coercion edge cases), ``where`` + selection and value propagation, and ``fma`` row/column broadcasting. + +**Internal** + +- New `bench_fma` and `bench_take` micro-benchmarks in + `scripts/bench_matrix.py`. The `examples/boids.py` demo migrates from + ``select`` to ``take``. + ## 2026-06-14 - Version 0.11.0 A behavior-dispatch release. `@when` becomes a **runtime decorator backed by a content-addressed marshalled-code registry** instead of a transpile-time diff --git a/CITATION.cff b/CITATION.cff index 78b8fe4..5fe4ba4 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -5,6 +5,6 @@ authors: given-names: "Matthew Alastair" orcid: "https://orcid.org/0000-0002-1019-8036" title: "bocpy" -version: 0.11.0 -date-released: 2026-06-14 +version: 0.12.0 +date-released: 2026-06-17 url: "https://github.com/microsoft/bocpy" \ No newline at end of file diff --git a/examples/boids.py b/examples/boids.py index f2da15c..6e0e830 100644 --- a/examples/boids.py +++ b/examples/boids.py @@ -313,8 +313,8 @@ def build_cell_data(self, positions: Matrix, velocities: Matrix, row: int, colum assert len(boids) > 0, "Invalid grid cell" - positions = positions.select(boids) - velocities = velocities.select(boids) + positions = positions.take(boids) + velocities = velocities.take(boids) return CellData(Cell(row, column), tuple(boids), Cown(positions), Cown(velocities)) def step(self, width: int, height: int): diff --git a/pyproject.toml b/pyproject.toml index 0b9553c..ad0d8f9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "bocpy" -version = "0.11.0" +version = "0.12.0" authors = [ {name = "bocpy Team", email="bocpy@microsoft.com"} ] diff --git a/scripts/bench_matrix.py b/scripts/bench_matrix.py index 46bae19..b101aa0 100644 --- a/scripts/bench_matrix.py +++ b/scripts/bench_matrix.py @@ -298,6 +298,48 @@ def bench_vecdot(results: list[dict[str, float]]) -> None: _print_row(results[-1]) +def bench_fma(results: list[dict[str, float]]) -> None: + """Bench fma(b, c) vs a*b + c. + + On the shipped x86-64 SSE2 build libc ``fma()`` is an unvectorisable + libcall, so ``a*b + c`` is expected to win on throughput there -- fma's + headline is accuracy (single rounding), not speed. A hardware-FMA host + (arm64, or x86-64 built with ``-mfma``) is where any throughput claim + must be measured. + """ + _print_section("fma(b, c) vs a*b + c (accuracy primitive; libcall on SSE2)") + + for rows, cols in [(1000, 3), (256, 256)]: + a = _make_matrix(rows, cols, seed=40) + b = _make_matrix(rows, cols, seed=41) + c = _make_matrix(rows, cols, seed=42) + label = f"{rows}x{cols}" + results.append(measure( + f"fma(b, c) matrix b,c {label}", + lambda a=a, b=b, c=c: a.fma(b, c), + )) + _print_row(results[-1]) + results.append(measure( + f"a*b + c matrix b,c {label}", + lambda a=a, b=b, c=c: a * b + c, + )) + _print_row(results[-1]) + + a = _make_matrix(256, 256, seed=40) + b = _make_matrix(256, 256, seed=41) + c = _make_matrix(256, 256, seed=42) + results.append(measure( + "fma(b, c, in_place) matrix b,c 256x256", + lambda a=a, b=b, c=c: a.fma(b, c, in_place=True), + )) + _print_row(results[-1]) + results.append(measure( + "fma(2.0, 1.0) scalar b,c 256x256", + lambda a=a: a.fma(2.0, 1.0), + )) + _print_row(results[-1]) + + def bench_cross(results: list[dict[str, float]]) -> None: """Bench cross across scalar and batch shapes.""" _print_section("cross at 1x3, 1000x2 (2D row batch), 1000x3 (3D row batch), 3x1000 (3D col batch)") @@ -624,22 +666,79 @@ def bench_transpose(results: list[dict[str, float]]) -> None: _print_row(results[-1]) -def bench_select(results: list[dict[str, float]]) -> None: - """Bench row and column gather across small and large index lists.""" - _print_section("select (row / column gather)") +def bench_take(results: list[dict[str, float]]) -> None: + """Bench row and column gather across wide, narrow, and column cases. + + The shared ``gather_axis`` helper copies each selected row with + ``memcpy`` (source and destination rows are both contiguous); the + wide row case (1000x100) is where that path dominates. The narrow-C + row case (1000x3) isolates the per-index unbox + bounds-check cost, + a larger fraction of the work when each copied row is short. Column + gather stays element-wise (inherently strided) and confirms no + regression there. + """ + _print_section("take (row / column gather)") shape = (1000, 100) m = _make_matrix(*shape, seed=108) row_idx = list(range(0, 1000, 10)) col_idx = list(range(0, 100, 2)) results.append(measure( - f"select rows (100/1000) shape={shape}", - lambda m=m, i=row_idx: m.select(i, 0), + f"take rows (100/1000) shape={shape}", + lambda m=m, i=row_idx: m.take(i, 0), + )) + _print_row(results[-1]) + results.append(measure( + f"take cols (50/100) shape={shape}", + lambda m=m, i=col_idx: m.take(i, 1), + )) + _print_row(results[-1]) + + narrow = _make_matrix(1000, 3, seed=51) + results.append(measure( + "take rows (narrow C) shape=(1000, 3) -> 100 rows", + lambda m=narrow, i=row_idx: m.take(i, 0), + )) + _print_row(results[-1]) + + +def bench_scatter(results: list[dict[str, float]]) -> None: + """Bench list-key scatter assignment (rows vs columns, fill vs matrix). + + Row scatter writes each selected destination row with a contiguous + ``memcpy`` (the headline write-side path); column scatter is strided. + The augmented ``m[[...]] += v`` case measures the gather -> in-place op + -> scatter triple pass against the plain single scatter; if that ratio + is large *and* the form is hot, that is the signal to revisit the + deferred fused single-pass kernel. + """ + _print_section("scatter (row / column assignment)") + + shape = (1000, 100) + row_idx = list(range(0, 1000, 10)) + col_idx = list(range(0, 100, 2)) + rows_rhs = _make_matrix(len(row_idx), 100, seed=110) + cols_rhs = _make_matrix(1000, len(col_idx), seed=111) + + m = _make_matrix(*shape, seed=112) + results.append(measure( + f"scatter rows = scalar shape={shape}", + lambda m=m, i=row_idx: m.__setitem__(i, 0.5), + )) + _print_row(results[-1]) + results.append(measure( + f"scatter rows = matrix shape={shape} (memcpy)", + lambda m=m, i=row_idx, v=rows_rhs: m.__setitem__(i, v), + )) + _print_row(results[-1]) + results.append(measure( + f"scatter cols = matrix shape={shape} (strided)", + lambda m=m, i=col_idx, v=cols_rhs: m.__setitem__((slice(None), i), v), )) _print_row(results[-1]) results.append(measure( - f"select cols (50/100) shape={shape}", - lambda m=m, i=col_idx: m.select(i, 1), + f"scatter rows += scalar shape={shape} (triple pass)", + lambda m=m, i=row_idx: m.__setitem__(i, m[i] + 0.5), )) _print_row(results[-1]) @@ -827,7 +926,8 @@ def main(argv: list[str] | None = None) -> int: bench_binary_arithmetic(results) bench_matmul(results) bench_transpose(results) - bench_select(results) + bench_take(results) + bench_scatter(results) bench_copy_clip_allclose(results) bench_construction(results) bench_factories(results) @@ -836,6 +936,7 @@ def main(argv: list[str] | None = None) -> int: bench_magnitude_squared(results) bench_negate(results) bench_vecdot(results) + bench_fma(results) bench_cross(results) bench_normalize(results) bench_perpendicular(results) diff --git a/sphinx/source/api.rst b/sphinx/source/api.rst index 80089d4..9487d72 100644 --- a/sphinx/source/api.rst +++ b/sphinx/source/api.rst @@ -125,7 +125,7 @@ Math .. autoclass:: Matrix :members: :undoc-members: - :special-members: __init__ + :special-members: __init__, __eq__, __ne__, __lt__, __le__, __gt__, __ge__, __getitem__, __setitem__ Messaging diff --git a/sphinx/source/conf.py b/sphinx/source/conf.py index 9cd8ab4..d1ea81f 100644 --- a/sphinx/source/conf.py +++ b/sphinx/source/conf.py @@ -14,7 +14,7 @@ project = 'bocpy' copyright = '2026, Microsoft' author = 'Microsoft' -release = '0.11.0' +release = '0.12.0' # -- General configuration --------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration diff --git a/src/bocpy/__init__.pyi b/src/bocpy/__init__.pyi index 86d21c3..87bbad8 100644 --- a/src/bocpy/__init__.pyi +++ b/src/bocpy/__init__.pyi @@ -229,6 +229,30 @@ class Matrix: :seealso: ``@`` / :func:`numpy.matmul` for matrix multiplication. """ + def fma(self, b: Union["Matrix", int, float], + c: Union["Matrix", int, float], + /, in_place: bool = False) -> "Matrix": + """Fused multiply-add: single-rounding ``self * b + c``. + + Uses C99 ``fma`` so the product is rounded once, unlike + ``self * b + c`` which rounds twice (the build sets + ``-ffp-contract=off`` for bit-reproducibility, so even an explicit + ``self * b + c`` never fuses). Results may differ from the + two-rounding form by up to half a ULP -- compare with + :meth:`allclose`, never ``==``. + + :param b: Multiplier: a same-shape :class:`Matrix`, a ``1x1`` + matrix, a ``1xN`` row vector or ``Mx1`` column vector that + broadcasts against ``self``, or a scalar. + :param c: Addend, with the same shape rules as *b*. + :param in_place: When ``True``, write into ``self`` and return it. + :return: ``self * b + c`` (``self`` itself when ``in_place=True``). + :raises ValueError: if *b* / *c* is a matrix whose shape neither + matches ``self`` nor broadcasts against it. + :raises TypeError: if *b* / *c* is neither a matrix nor a real + number. + """ + def cross(self, other: "Matrix", axis: Optional[int] = None) -> Union[float, "Matrix"]: """2D or 3D cross product against another vector or batch. @@ -265,7 +289,7 @@ class Matrix: through this method). :return: A float for ``1x2`` / ``2x1`` inputs; otherwise a :class:`Matrix`. - :raises NotImplementedError: on incompatible shapes or + :raises ValueError: on incompatible shapes or mismatched batch sizes. """ @@ -302,7 +326,7 @@ class Matrix: When ``False`` (the default), return a new :class:`Matrix` with the rotated vectors. :return: A :class:`Matrix` (``self`` when ``in_place=True``). - :raises NotImplementedError: on any shape that is not a 2D vector + :raises ValueError: on any shape that is not a 2D vector or a ``Nx2`` / ``2xN`` batch. """ @@ -319,7 +343,7 @@ class Matrix: on unambiguous shapes. :return: A float for a single 2D vector input, otherwise a :class:`Matrix` of per-vector angles. - :raises NotImplementedError: on any shape that is not a 2D vector + :raises ValueError: on any shape that is not a 2D vector or a ``Nx2`` / ``2xN`` batch. """ @@ -397,15 +421,138 @@ class Matrix: :param in_place: When ``True``, mutate ``self`` and return it. """ - def clip(self, min_or_maxval: float, maxval: Optional[float] = None) -> "Matrix": - """Clip every element to a given range. + def sqrt(self, in_place: bool = False) -> "Matrix": + """Take the square root of every element. - :param min_or_maxval: If *maxval* is provided, this is the minimum - clipping value. Otherwise, this is the maximum and the minimum - defaults to zero. - :param maxval: The maximum clipping value, or ``None`` to treat - *min_or_maxval* as the maximum. + Negative elements yield ``NaN`` (no exception is raised), matching + :func:`numpy.sqrt`. + + :param in_place: When ``True``, mutate ``self`` and return it. + """ + + def less(self, other: Union["Matrix", int, float, + Sequence[Union[int, float]]]) -> "Matrix": + """Element-wise ``self < other`` as a 0/1 mask matrix. + + Distinct from the ``<`` operator, which returns a single + :class:`bool` (see :meth:`__lt__`). + + :param other: A same-shape matrix, a scalar (including ``bool``), a + ``1x1`` matrix, a row/column vector that broadcasts (same rules + as arithmetic), or a list/tuple of numbers. + :return: A new :class:`Matrix` of ``1.0``/``0.0``. NaN comparisons + yield ``0.0``. + :raises ValueError: on a non-broadcastable shape or an empty + list/tuple operand. + :raises TypeError: if *other* is a list/tuple holding a non-number, + or is not a matrix, scalar, or list/tuple. + """ + + def less_equal(self, other: Union["Matrix", int, float, + Sequence[Union[int, float]]]) -> "Matrix": + """Element-wise ``self <= other`` as a 0/1 mask matrix. + + Distinct from the ``<=`` operator, which returns a single + :class:`bool` (see :meth:`__le__`). + + :param other: A same-shape matrix, a scalar (including ``bool``), a + ``1x1`` matrix, a row/column vector that broadcasts, or a + list/tuple of numbers. + :return: A new :class:`Matrix` of ``1.0``/``0.0``. NaN comparisons + yield ``0.0``. + :raises ValueError: on a non-broadcastable shape or an empty + list/tuple operand. + :raises TypeError: on a non-number list/tuple element or an + unsupported operand type (see :meth:`less`). + """ + + def greater(self, other: Union["Matrix", int, float, + Sequence[Union[int, float]]]) -> "Matrix": + """Element-wise ``self > other`` as a 0/1 mask matrix. + + Distinct from the ``>`` operator, which returns a single + :class:`bool` (see :meth:`__gt__`). + + :param other: A same-shape matrix, a scalar (including ``bool``), a + ``1x1`` matrix, a row/column vector that broadcasts, or a + list/tuple of numbers. + :return: A new :class:`Matrix` of ``1.0``/``0.0``. NaN comparisons + yield ``0.0``. + :raises ValueError: on a non-broadcastable shape or an empty + list/tuple operand. + :raises TypeError: on a non-number list/tuple element or an + unsupported operand type (see :meth:`less`). + """ + + def greater_equal(self, other: Union["Matrix", int, float, + Sequence[Union[int, float]]]) -> "Matrix": + """Element-wise ``self >= other`` as a 0/1 mask matrix. + + Distinct from the ``>=`` operator, which returns a single + :class:`bool` (see :meth:`__ge__`). + + :param other: A same-shape matrix, a scalar (including ``bool``), a + ``1x1`` matrix, a row/column vector that broadcasts, or a + list/tuple of numbers. + :return: A new :class:`Matrix` of ``1.0``/``0.0``. NaN comparisons + yield ``0.0``. + :raises ValueError: on a non-broadcastable shape or an empty + list/tuple operand. + :raises TypeError: on a non-number list/tuple element or an + unsupported operand type (see :meth:`less`). + """ + + def equal(self, other: Union["Matrix", int, float, + Sequence[Union[int, float]]]) -> "Matrix": + """Element-wise ``self == other`` as a 0/1 mask matrix. + + Distinct from the ``==`` operator, which returns a single + :class:`bool` (see :meth:`__eq__`). + + :param other: A same-shape matrix, a scalar (including ``bool``), a + ``1x1`` matrix, a row/column vector that broadcasts, or a + list/tuple of numbers. + :return: A new :class:`Matrix` of ``1.0``/``0.0``. NaN comparisons + yield ``0.0``. + :raises ValueError: on a non-broadcastable shape or an empty + list/tuple operand. + :raises TypeError: on a non-number list/tuple element or an + unsupported operand type (see :meth:`less`). + """ + + def not_equal(self, other: Union["Matrix", int, float, + Sequence[Union[int, float]]]) -> "Matrix": + """Element-wise ``self != other`` as a 0/1 mask matrix. + + Distinct from the ``!=`` operator, which returns a single + :class:`bool` (see :meth:`__ne__`). + + :param other: A same-shape matrix, a scalar (including ``bool``), a + ``1x1`` matrix, a row/column vector that broadcasts, or a + list/tuple of numbers. + :return: A new :class:`Matrix` of ``1.0``/``0.0``. NaN comparisons + yield ``1.0``. + :raises ValueError: on a non-broadcastable shape or an empty + list/tuple operand. + :raises TypeError: on a non-number list/tuple element or an + unsupported operand type (see :meth:`less`). + """ + + def clip(self, min: Optional[float] = None, + max: Optional[float] = None) -> "Matrix": + """Clamp every element to ``[min, max]``. + + The first argument is the lower bound and the second the upper + bound. Either may be ``None`` (or omitted) to leave that side + unbounded: ``m.clip(min=0.0)`` clamps only below and + ``m.clip(max=255.0)`` only above. + + :param min: Lower clipping bound, or ``None`` for no lower bound. + :param max: Upper clipping bound, or ``None`` for no upper bound. :return: A new clipped :class:`Matrix`. + :raises ValueError: if both *min* and *max* are ``None``. + :raises AssertionError: if both bounds are given and *max* < *min*. + :raises TypeError: if a given bound is not a real number. """ def copy(self) -> "Matrix": @@ -419,11 +566,64 @@ class Matrix: object overhead. The current interpreter must own the matrix. """ - def select(self, indices: Union[list[int], tuple[int]], axis=0): + def take(self, indices: Union[list[int], tuple[int]], axis=0) -> "Matrix": """Return a new matrix containing only the selected rows or columns. - :param indices: The row or column indices to select. - :param axis: ``0`` to select rows, ``1`` to select columns. + *indices* is a 1-D list or tuple of ints. Negative indices count + from the end. Duplicate indices repeat the corresponding row or + column. Equivalent to ``m[indices]`` (rows) and ``m[:, indices]`` + (columns). A ``bool`` element is treated as the integer ``0``/``1``. + + :param indices: The row or column indices to take. Must be a + non-empty list or tuple of ints. + :param axis: ``0`` to take rows, ``1`` to take columns. + :raises IndexError: if an index is out of range, or if *indices* + is empty. + :raises KeyError: if *axis* is not ``0`` or ``1``. + :raises TypeError: if an index is not an int. + :raises OverflowError: if an index exceeds the platform word size. + """ + + def put(self, indices: Union[list[int], tuple[int]], + value: Union[int, float, "Matrix"], axis=0, + accumulate: bool = False) -> "Matrix": + """Assign *value* into the selected rows or columns in place. + + The write-side counterpart of :meth:`take`. *value* may be a + scalar (a real number or a ``1x1`` matrix, broadcast over the + selection) or a matrix whose shape matches the selection exactly + (``len(indices)`` rows by this matrix's column count for a row + assignment; this matrix's row count by ``len(indices)`` columns for + a column assignment). Equivalent to ``m[indices] = value`` (rows) + and ``m[:, indices] = value`` (columns). + + All indices and the *value* shape are validated before any element + is written, so a rejected call leaves the matrix unchanged. Negative + indices count from the end. A ``bool`` index element is treated as + the integer ``0``/``1``. + + With ``accumulate=False`` (the default) duplicate indices follow + last-write-wins. With ``accumulate=True`` the values are *added* + into the selection, so duplicate indices fold additively + (``m.put([0, 0], v, accumulate=True)`` adds twice into row 0). This + is the only way to fold duplicates: ``m[indices] += value`` desugars + to gather/iadd/scatter and so collapses to last-write-wins. + + :param indices: The row or column indices to assign. Must be a + non-empty list or tuple of ints. + :param value: The scalar or matrix to assign into the selection. + :param axis: ``0`` to assign rows, ``1`` to assign columns. + :param accumulate: When ``True``, add into the selection instead of + overwriting, so duplicate indices accumulate. + :return: ``self`` (to allow chaining). + :raises IndexError: if an index is out of range, or if *indices* + is empty. + :raises KeyError: if *axis* is not ``0`` or ``1``. + :raises ValueError: if a matrix *value* shape does not match the + selection shape. + :raises TypeError: if *value* is neither a real number nor a matrix, + or if an index is not an int. + :raises OverflowError: if an index exceeds the platform word size. """ def __add__(self, other: Union["Matrix", int, float]) -> "Matrix": @@ -459,16 +659,154 @@ class Matrix: def __len__(self) -> int: """Return the number of rows.""" - def __getitem__(self, key: Union[int, tuple[int, int]]) -> Union["Matrix", float]: - """Retrieve a row, element, or sub-matrix by index or slice.""" + def __getitem__(self, key: Union[int, slice, tuple, list[int]]) -> Union["Matrix", float]: + """Retrieve a row, element, sub-matrix, or fancy-index gather. + + A list key gathers rows or columns: ``m[[r0, r1]]`` and + ``m[[r0, r1], :]`` select rows, and ``m[:, [c0, c1]]`` selects + columns. A list key always returns a :class:`Matrix` (even a + single-element list such as ``m[[0]]``), whereas ``m[0]`` and + ``m[0, 0]`` return a Python ``float``. Negative indices count from + the end; duplicates repeat the row or column; an out-of-range index + raises :class:`IndexError` and an empty list raises + :class:`IndexError`. + + Column gather requires the bare ``:`` row selector: ``m[0:R, [c]]`` + (a full *range* rather than ``:``) raises + :class:`IndexError`. Paired lists (``m[[r], [c]]``), + list-with-int (``m[[r], 0]``), and a list paired with a non-full + slice all raise :class:`IndexError`; use :meth:`take` + for those. A ``bool`` element in a list is treated as ``0``/``1``. + + :raises IndexError: if a list index is out of range, the list + is empty, or for unsupported list/slice combinations. + """ - def __setitem__(self, key: Union[int, tuple[int, int]], value: Union[int, - float, "Matrix", Sequence[Union[int, float]]]): - """Set a row, element, or sub-matrix by index or slice.""" + def __setitem__(self, key: Union[int, slice, tuple, list[int]], + value: Union[int, float, "Matrix", + Sequence[Union[int, float]]]): + """Set a row, element, sub-matrix, or fancy-index scatter. + + A list key scatters into rows or columns, mirroring the + :meth:`__getitem__` gather shapes: ``m[[r0, r1]] = v`` and + ``m[[r0, r1], :] = v`` assign rows, and ``m[:, [c0, c1]] = v`` + assigns columns. The right-hand side may be a scalar (a real number + or a ``1x1`` matrix, broadcast over the selection) or a matrix whose + shape matches the selection exactly (``count`` rows by the receiver's + column count for a row scatter; the receiver's row count by ``count`` + columns for a column scatter). + + All indices and the RHS shape are validated *before* any element is + written, so a rejected assignment leaves the matrix unchanged (no + partial writes). Duplicate indices follow last-write-wins, **not** + accumulation — ``m[[0, 0]] += v`` increments row 0 once, not twice; + use :meth:`put` with ``accumulate=True`` to fold duplicates. The + augmented forms (``+=``, ``-=``, ``*=``, ``/=``) desugar to a + gather, an in-place op, and a scatter. + + Column scatter requires the bare ``:`` row selector: ``m[0:R, [c]]`` + raises :class:`IndexError`. Paired lists (``m[[r], [c]]``), + list-with-int (``m[[r], 0]``), and a list paired with a non-full + slice all raise :class:`IndexError`. A ``bool`` element in + a list is treated as ``0``/``1``. + + :raises IndexError: if a list index is out of range, the list + is empty, or for unsupported list/slice combinations. + :raises ValueError: if a matrix RHS shape does not match the + selection shape. + :raises TypeError: if the RHS is neither a real number nor a matrix. + """ def __iter__(self) -> Iterator[Union[float, "Matrix"]]: """Iterate over rows of the matrix.""" + def __lt__(self, other: Union["Matrix", int, float, + Sequence[Union[int, float]]]) -> bool: + """Lexicographic ``self < other`` returning a single :class:`bool`. + + Compares element by element in row-major order, like a list or + tuple: the first element where a strict ordering holds decides. + *other* must be a same-shape matrix, a scalar (a scalar + broadcasts to ``self``'s shape), or a list/tuple of numbers (coerced + to a ``1xN`` row matrix). A ``1x1`` matrix is **not** treated + as a scalar here and only compares against another ``1x1``. A + ``NaN`` element never decides the ordering; comparison continues + past it, so two all-``NaN`` matrices compare neither ``<`` nor ``>``. + + For an element-wise 0/1 mask instead, use :meth:`less`. + + :raises ValueError: if *other* is a matrix (or coerced sequence) of + a different shape, or an empty list/tuple. + :raises TypeError: if *other* is a list/tuple holding a non-number. + """ + + def __le__(self, other: Union["Matrix", int, float, + Sequence[Union[int, float]]]) -> bool: + """Lexicographic ``self <= other`` returning a single :class:`bool`. + + See :meth:`__lt__` for the comparison rules. For an element-wise + 0/1 mask instead, use :meth:`less_equal`. + + :raises ValueError: if *other* is a matrix (or coerced sequence) of + a different shape, or an empty list/tuple. + :raises TypeError: if *other* is a list/tuple holding a non-number. + """ + + def __gt__(self, other: Union["Matrix", int, float, + Sequence[Union[int, float]]]) -> bool: + """Lexicographic ``self > other`` returning a single :class:`bool`. + + See :meth:`__lt__` for the comparison rules. For an element-wise + 0/1 mask instead, use :meth:`greater`. + + :raises ValueError: if *other* is a matrix (or coerced sequence) of + a different shape, or an empty list/tuple. + :raises TypeError: if *other* is a list/tuple holding a non-number. + """ + + def __ge__(self, other: Union["Matrix", int, float, + Sequence[Union[int, float]]]) -> bool: + """Lexicographic ``self >= other`` returning a single :class:`bool`. + + See :meth:`__lt__` for the comparison rules. For an element-wise + 0/1 mask instead, use :meth:`greater_equal`. + + :raises ValueError: if *other* is a matrix (or coerced sequence) of + a different shape, or an empty list/tuple. + :raises TypeError: if *other* is a list/tuple holding a non-number. + """ + + def __eq__(self, other: object) -> bool: + """Lexicographic ``self == other`` returning a single :class:`bool`. + + ``True`` only when *other* is a same-shape matrix equal element by + element, a list/tuple of numbers (coerced to a ``1xN`` row matrix) + equal element by element, or a scalar that every element equals. + Total over non-equality: a shape mismatch, a non-matrix / non-scalar + / non-list-tuple operand, or a list/tuple that cannot be coerced + (empty, or holding a non-number) all return ``False`` rather than + raising, so ``matrix in some_list`` still works. A ``NaN`` element + never decides the comparison (it is skipped), so an all-``NaN`` + matrix compares ``==`` equal to itself; this diverges from the + :meth:`equal` mask, where ``NaN == x`` is ``0.0``. Defining value + equality makes :class:`Matrix` unhashable (so it cannot be a dict + key or set member), which is correct for a mutable type. For an + element-wise 0/1 mask instead, use :meth:`equal`. + + :raises RuntimeError: if *other* is a matrix owned by a different + interpreter. + """ + + def __ne__(self, other: object) -> bool: + """Lexicographic ``self != other`` returning a single :class:`bool`. + + The (total) negation of :meth:`__eq__`. For an element-wise + 0/1 mask instead, use :meth:`not_equal`. + + :raises RuntimeError: if *other* is a matrix owned by a different + interpreter. + """ + @classmethod def allclose(cls, lhs: "Matrix", rhs: "Matrix", rtol: float = 1e-05, atol: float = 1e-08, equal_nan: bool = False): """Check whether two matrices are element-wise equal within a tolerance. @@ -483,6 +821,32 @@ class Matrix: :return: ``True`` if every element pair satisfies the tolerance. """ + @classmethod + def where(cls, mask: "Matrix", + a: Union["Matrix", int, float, Sequence[Union[int, float]]], + b: Union["Matrix", int, float, + Sequence[Union[int, float]]]) -> "Matrix": + """Select element-wise from *a* or *b* on a truthy mask. + + Returns a fresh matrix taking *a* where the corresponding *mask* + element is non-zero and *b* elsewhere, like :func:`numpy.where`. + + :param mask: A :class:`Matrix` whose non-zero elements select *a*. + ``NaN`` mask elements count as non-zero and select *a*. + :param a: A scalar (including ``bool``), a list/tuple of numbers + (coerced to a ``1xN`` row matrix), or a :class:`Matrix` matching + *mask*'s shape. + :param b: A scalar (including ``bool``), a list/tuple of numbers, or + a :class:`Matrix` matching *mask*'s shape. + :return: A new :class:`Matrix` with *mask*'s shape. + :raises TypeError: if *a* or *b* is neither a matrix, a scalar, nor + a list/tuple of numbers, or is a list/tuple holding a non-number. + :raises ValueError: if a matrix (or coerced sequence) operand's + shape differs from *mask*'s shape, or is an empty list/tuple. A + ``1x1`` matrix is treated as a matrix (not a scalar) and so must + match the mask shape. + """ + @classmethod def zeros(cls, size: tuple[int, int]): """Create a matrix filled with zeros. diff --git a/src/bocpy/_math.c b/src/bocpy/_math.c index a1852cc..e869834 100644 --- a/src/bocpy/_math.c +++ b/src/bocpy/_math.c @@ -33,6 +33,32 @@ #define BOC_CANARY_NOINLINE #endif +/* Runtime CPU dispatch for the fma kernel (GNU function multiversioning). + target_clones emits a hardware-FMA body plus a portable default body and + binds the right one at load time via an STT_GNU_IFUNC resolver. This lets + the shipped x86-64 wheel use vfmadd on capable CPUs while the SSE2 + baseline still gets a correct (libcall) body. + + Gated on __GLIBC__ deliberately: this is the ONLY platform where the + mechanism is safe, confirmed by compiling the kernel inside the actual + cibuildwheel default containers (2026-06-16): + - manylinux2014 (glibc): IFUNC resolves, correct result. + - musllinux_1_2 (musl): HARD COMPILE ERROR — "the call requires + 'ifunc', which is not supported by this + target". An unguarded attribute would break + the entire musllinux wheel leg. + macOS (Mach-O) and Windows (PE/COFF) also lack IFUNC; __GLIBC__ excludes + all three (musl, macOS, Windows) in a single test. On glibc the + include above defines __GLIBC__ via . Canary builds opt out + so per-helper disassembly stays a single stable body. */ +#if defined(__x86_64__) && defined(__GLIBC__) && \ + (defined(__GNUC__) || defined(__clang__)) && \ + !defined(BOC_CANARY_NOINLINE_ON) +#define BOC_FMA_MULTIVERSION __attribute__((target_clones("fma", "default"))) +#else +#define BOC_FMA_MULTIVERSION +#endif + /// @brief Underlying C-based matrix implementation typedef struct boc_matrix_impl { /// @brief The raw double values of the matrix @@ -234,7 +260,13 @@ enum BinaryOps { RSubtract = 1002, Multiply = 1003, Divide = 1004, - RDivide = 1005 + RDivide = 1005, + Less = 1006, + LessEqual = 1007, + Greater = 1008, + GreaterEqual = 1009, + Equal = 1010, + NotEqual = 1011 }; /* -------------------------------------------------------------------------- @@ -272,7 +304,13 @@ enum BinaryOps { X(RSubtract, rsubtract, rhs - lhs) \ X(Multiply, multiply, lhs *rhs) \ X(Divide, divide, lhs / rhs) \ - X(RDivide, rdivide, rhs / lhs) + X(RDivide, rdivide, rhs / lhs) \ + X(Less, less, (double)(lhs < rhs)) \ + X(LessEqual, less_equal, (double)(lhs <= rhs)) \ + X(Greater, greater, (double)(lhs > rhs)) \ + X(GreaterEqual, greater_equal, (double)(lhs >= rhs)) \ + X(Equal, equal, (double)(lhs == rhs)) \ + X(NotEqual, not_equal, (double)(lhs != rhs)) #define DEFINE_BINARY_EWISE(ENUM, STAMP, EXPR) \ BOC_CANARY_NOINLINE \ @@ -780,7 +818,8 @@ static void dispatch_agg_columnwise(matrix_impl *m, enum AggregateOps op, X(Floor, floor, floor(v)) \ X(Round, round, nearbyint(v)) \ X(Negate, negate, -v) \ - X(Abs, abs, fabs(v)) + X(Abs, abs, fabs(v)) \ + X(Sqrt, sqrt, sqrt(v)) #define DEFINE_UNARY(ENUM, STAMP, EXPR) \ BOC_CANARY_NOINLINE \ @@ -988,6 +1027,12 @@ static matrix_impl *impl_new_from_sequence(PyObject *sequence, bool as_column) { } Py_ssize_t size = PySequence_Fast_GET_SIZE(fast); + if (size == 0) { + Py_DECREF(fast); + PyErr_SetString(PyExc_ValueError, + "cannot coerce an empty list/tuple to a matrix"); + return NULL; + } matrix_impl *impl; if (as_column) { impl = impl_new((size_t)size, 1); @@ -1316,7 +1361,7 @@ static int parse_validate_normalise_axis(PyObject *axis_obj, AxisArg *out) { } int ax = out->axis; if (ax != 0 && ax != 1 && ax != -1 && ax != -2) { - PyErr_SetString(PyExc_NotImplementedError, "axis must be -2, -1, 0, or 1"); + PyErr_SetString(PyExc_ValueError, "axis must be -2, -1, 0, or 1"); return -1; } if (ax == -1) { @@ -1815,7 +1860,7 @@ static PyObject *Matrix_vecdot(PyObject *op, PyObject *args, PyObject *kwds) { (rhs->rows == 1 || rhs->columns == 1) && lhs->size == rhs->size) { shape = BCAST_NONE; } else { - PyErr_Format(PyExc_NotImplementedError, + PyErr_Format(PyExc_ValueError, "vecdot: lhs %zux%zu incompatible with rhs %zux%zu", lhs->rows, lhs->columns, rhs->rows, rhs->columns); goto done; @@ -1849,6 +1894,190 @@ static PyObject *Matrix_vecdot(PyObject *op, PyObject *args, PyObject *kwds) { return result; } +/// @brief A classified fma operand: either a scalar or a same-shape matrix. +/// @details ``full`` is an INCREF'd matrix with the same shape as the +/// receiver, and the caller must IMPL_DECREF it; when ``full`` is +/// NULL the operand is the scalar in ``scalar``. +typedef struct { + double scalar; + matrix_impl *full; +} fma_operand; + +/// @brief Materialise a row- or column-vector broadcast into a full buffer. +/// @details Expands a ``1xN`` row vector (``is_row``) or an ``Mx1`` column +/// vector into a fresh contiguous ``rows`` x ``columns`` matrix so +/// the contiguous fma kernel keeps its unit stride (and hardware +/// FMA). The returned impl has refcount 0; the caller INCREFs it. +static matrix_impl *impl_broadcast_vector(const matrix_impl *vec, size_t rows, + size_t columns, bool is_row) { + matrix_impl *out = impl_new(rows, columns); + if (out == NULL) { + return NULL; + } + double *dst = out->data; + const double *v = vec->data; + if (is_row) { + for (size_t r = 0; r < rows; ++r) { + for (size_t c = 0; c < columns; ++c) { + *dst++ = v[c]; + } + } + } else { + for (size_t r = 0; r < rows; ++r) { + for (size_t c = 0; c < columns; ++c) { + *dst++ = v[r]; + } + } + } + return out; +} + +/// @brief Classify an fma operand as a scalar or a full same-shape matrix. +/// @details A Matrix whose size is 1 is treated as a scalar so a ``1x1`` +/// broadcasts like a number (the in-house scalar rule used +/// elsewhere); a same-shape Matrix is INCREF'd into ``out->full``. +/// A ``1xN`` row vector (matching self's columns) or an ``Mx1`` +/// column vector (matching self's rows) is materialised into a fresh +/// same-shape buffer stored in ``out->full`` so the contiguous kernel +/// keeps its unit stride; any other matrix shape raises ``ValueError`` +/// naming both shapes. Any real number is accepted as a scalar. +/// @param op The operand object (``b`` or ``c``). +/// @param self_impl The receiver's impl, used for the shape check. +/// @param what Operand name used in error messages. +/// @param out Output operand descriptor (zeroed before classification). +/// @return 0 on success; -1 with an exception set on failure. +static int classify_fma_operand(PyObject *op, const matrix_impl *self_impl, + const char *what, fma_operand *out) { + out->scalar = 0.0; + out->full = NULL; + + if (Py_TYPE(op) == LOCAL_STATE->matrix_type) { + matrix_impl *impl = ((MatrixObject *)op)->impl; + if (!impl_check_acquired(impl, true)) { + return -1; + } + if (impl->size == 1) { + out->scalar = impl->data[0]; + return 0; + } + if (impl->rows == self_impl->rows && impl->columns == self_impl->columns) { + IMPL_INCREF(impl); + out->full = impl; + return 0; + } + if (impl->rows == 1 && impl->columns == self_impl->columns) { + matrix_impl *full = impl_broadcast_vector(impl, self_impl->rows, + self_impl->columns, true); + if (full == NULL) { + return -1; + } + IMPL_INCREF(full); + out->full = full; + return 0; + } + if (impl->columns == 1 && impl->rows == self_impl->rows) { + matrix_impl *full = impl_broadcast_vector(impl, self_impl->rows, + self_impl->columns, false); + if (full == NULL) { + return -1; + } + IMPL_INCREF(full); + out->full = full; + return 0; + } + PyErr_Format(PyExc_ValueError, + "fma: %s shape %zux%zu incompatible with self %zux%zu", what, + impl->rows, impl->columns, self_impl->rows, + self_impl->columns); + return -1; + } + + if (unwrap_double(op, &out->scalar)) { + return 0; + } + if (PyErr_Occurred()) { + return -1; + } + + PyErr_Format(PyExc_TypeError, "fma: %s must be a Matrix or a real number", + what); + return -1; +} + +/// @brief Fused multiply-add kernel: ``out = a*b + c`` with one rounding. +/// @details ``b_step`` / ``c_step`` are 1 for a full same-shape operand and +/// 0 to repeat a scalar (the pointer then aims at a single local +/// ``double``). Each output cell reads only its own index, so the +/// in-place case (``out == a``) and ``b`` / ``c`` aliasing ``a`` +/// are all safe. ``fma()`` rounds the product once, unlike +/// ``a*b + c`` which rounds twice under ``-ffp-contract=off``. +/// BOC_FMA_MULTIVERSION clones this kernel for hardware FMA on +/// glibc x86-64 (see the macro definition for why glibc-only). +BOC_CANARY_NOINLINE +BOC_FMA_MULTIVERSION +static void impl_fma(const double *a, const double *b, size_t b_step, + const double *c, size_t c_step, double *out, size_t n) { + for (size_t i = 0; i < n; ++i, ++a, b += b_step, c += c_step, ++out) { + *out = fma(*a, *b, *c); + } +} + +/// @brief Fused multiply-add: single-rounding ``self*b + c``. +/// @details ``b`` and ``c`` may each be a same-shape matrix, a ``1x1`` +/// matrix, a ``1xN`` row vector or ``Mx1`` column vector that +/// broadcasts against ``self``, or a scalar; any other matrix shape +/// raises ``ValueError``. Both operands are validated before any +/// allocation, so a rejected operand leaves ``self`` untouched. +/// With ``in_place=True`` the result is written into ``self`` and +/// ``self`` is returned. +static PyObject *Matrix_fma(PyObject *op, PyObject *args, PyObject *kwds) { + MatrixObject *self = (MatrixObject *)op; + PyObject *b_op = NULL; + PyObject *c_op = NULL; + int in_place = 0; + PyObject *out_op = NULL; + fma_operand bop = {0.0, NULL}; + fma_operand cop = {0.0, NULL}; + + static char *kwlist[] = {"", "", "in_place", NULL}; + if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|p", kwlist, &b_op, &c_op, + &in_place)) { + return NULL; + } + + if (!impl_check_acquired(self->impl, true)) { + return NULL; + } + + if (classify_fma_operand(b_op, self->impl, "b", &bop) < 0) { + goto done; + } + if (classify_fma_operand(c_op, self->impl, "c", &cop) < 0) { + goto done; + } + + matrix_impl *out = set_output(op, &out_op, in_place); + if (out == NULL) { + out_op = NULL; + goto done; + } + + double b_scalar = bop.scalar; + double c_scalar = cop.scalar; + const double *b_ptr = bop.full != NULL ? bop.full->data : &b_scalar; + const double *c_ptr = cop.full != NULL ? cop.full->data : &c_scalar; + size_t b_step = bop.full != NULL ? 1 : 0; + size_t c_step = cop.full != NULL ? 1 : 0; + + impl_fma(self->impl->data, b_ptr, b_step, c_ptr, c_step, out->data, + self->impl->size); + +done: + IMPL_DECREF(bop.full); + IMPL_DECREF(cop.full); + return out_op; +} + /// @brief Classify a matrix as a 2D/3D cross-product operand or batch. /// @details ``has_axis`` / ``explicit_axis`` carry an optional caller- /// supplied axis (already normalised to 0 or 1 by @@ -1990,14 +2219,14 @@ static PyObject *Matrix_cross(PyObject *op, PyObject *args, PyObject *kwds) { enum CrossAxis flavor = classify_cross_axis(lhs, axis.has_axis, axis.axis); if (flavor == CROSS_INVALID) { PyErr_SetString( - PyExc_NotImplementedError, + PyExc_ValueError, "cross requires a 2D or 3D vector or Nx2 or 2xN or Nx3 or 3xN matrix"); goto done; } if (flavor == CROSS_SCALAR_2D_1x2 || flavor == CROSS_SCALAR_2D_2x1) { if (rhs->size != 2) { - PyErr_Format(PyExc_NotImplementedError, + PyErr_Format(PyExc_ValueError, "cross: 2D vector lhs %zux%zu incompatible with rhs %zux%zu", lhs->rows, lhs->columns, rhs->rows, rhs->columns); goto done; @@ -2009,7 +2238,7 @@ static PyObject *Matrix_cross(PyObject *op, PyObject *args, PyObject *kwds) { } if (flavor == CROSS_SCALAR_3D_1x3 || flavor == CROSS_SCALAR_3D_3x1) { if (rhs->size != 3) { - PyErr_Format(PyExc_NotImplementedError, + PyErr_Format(PyExc_ValueError, "cross: 3D vector lhs %zux%zu incompatible with rhs %zux%zu", lhs->rows, lhs->columns, rhs->rows, rhs->columns); goto done; @@ -2034,7 +2263,7 @@ static PyObject *Matrix_cross(PyObject *op, PyObject *args, PyObject *kwds) { const bool broadcast = (rhs->size == 2 && (rhs->rows == 1 || rhs->columns == 1)); if (!same_shape && !broadcast) { - PyErr_Format(PyExc_NotImplementedError, + PyErr_Format(PyExc_ValueError, "cross: Nx2 batch lhs %zux%zu incompatible with rhs %zux%zu", lhs->rows, lhs->columns, rhs->rows, rhs->columns); goto done; @@ -2074,7 +2303,7 @@ static PyObject *Matrix_cross(PyObject *op, PyObject *args, PyObject *kwds) { const bool broadcast = (rhs->size == 2 && (rhs->rows == 1 || rhs->columns == 1)); if (!same_shape && !broadcast) { - PyErr_Format(PyExc_NotImplementedError, + PyErr_Format(PyExc_ValueError, "cross: 2xN batch lhs %zux%zu incompatible with rhs %zux%zu", lhs->rows, lhs->columns, rhs->rows, rhs->columns); goto done; @@ -2110,7 +2339,7 @@ static PyObject *Matrix_cross(PyObject *op, PyObject *args, PyObject *kwds) { const bool broadcast = (rhs->size == 3 && (rhs->rows == 1 || rhs->columns == 1)); if (!same_shape && !broadcast) { - PyErr_Format(PyExc_NotImplementedError, + PyErr_Format(PyExc_ValueError, "cross: Nx3 batch lhs %zux%zu incompatible with rhs %zux%zu", lhs->rows, lhs->columns, rhs->rows, rhs->columns); goto done; @@ -2157,7 +2386,7 @@ static PyObject *Matrix_cross(PyObject *op, PyObject *args, PyObject *kwds) { const bool broadcast = (rhs->size == 3 && (rhs->rows == 1 || rhs->columns == 1)); if (!same_shape && !broadcast) { - PyErr_Format(PyExc_NotImplementedError, + PyErr_Format(PyExc_ValueError, "cross: 3xN batch lhs %zux%zu incompatible with rhs %zux%zu", lhs->rows, lhs->columns, rhs->rows, rhs->columns); goto done; @@ -2427,7 +2656,7 @@ static PyObject *Matrix_perpendicular(PyObject *op, PyObject *args, enum Vec2Axis flavor = classify_vec2_axis(self->impl, axis.has_axis, axis.axis); if (flavor == VEC2_INVALID) { - PyErr_SetString(PyExc_NotImplementedError, + PyErr_SetString(PyExc_ValueError, "perpendicular requires a 2D vector or Nx2 or 2xN matrix"); return NULL; } @@ -2470,7 +2699,7 @@ static PyObject *Matrix_angle(PyObject *op, PyObject *args, PyObject *kwds) { enum Vec2Axis flavor = classify_vec2_axis(self->impl, axis.has_axis, axis.axis); if (flavor == VEC2_INVALID) { - PyErr_SetString(PyExc_NotImplementedError, + PyErr_SetString(PyExc_ValueError, "angle requires a 2D vector or Nx2 or 2xN matrix"); return NULL; } @@ -2544,7 +2773,13 @@ static PyObject *Matrix_angle(PyObject *op, PyObject *args, PyObject *kwds) { BOC_UNARY_OPS(X) #undef X -static PyObject *Matrix_clip(PyObject *op, PyObject *args) { +/// @brief Clamp every element to ``[min, max]``; either bound may be omitted +/// @details The first positional argument is the lower bound, the second the +/// upper. An omitted bound (``None``) is realised as +/// ``-INFINITY`` / ``+INFINITY`` so the inner loop stays branch-free; +/// NaN elements pass through unchanged. Raises ``ValueError`` if both +/// bounds are omitted and ``AssertionError`` if ``max < min``. +static PyObject *Matrix_clip(PyObject *op, PyObject *args, PyObject *kwds) { MatrixObject *self = (MatrixObject *)op; matrix_impl *impl = self->impl; @@ -2552,34 +2787,34 @@ static PyObject *Matrix_clip(PyObject *op, PyObject *args) { return NULL; } - PyObject *minval_op = NULL; - PyObject *maxval_op = NULL; + PyObject *minval_op = Py_None; + PyObject *maxval_op = Py_None; + static char *kwlist[] = {"min", "max", NULL}; - if (!PyArg_ParseTuple(args, "O|O", &minval_op, &maxval_op)) { + if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO", kwlist, &minval_op, + &maxval_op)) { return NULL; } - double minval; - double maxval; - if (maxval_op == NULL) { - minval = 0; - if (!unwrap_double(minval_op, &maxval)) { - PyErr_SetString(PyExc_TypeError, "Expected a number"); - return NULL; - } - } else { - if (!unwrap_double(minval_op, &minval)) { - PyErr_SetString(PyExc_TypeError, "Expected a number"); - return NULL; - } + bool has_min = minval_op != Py_None; + bool has_max = maxval_op != Py_None; + if (!has_min && !has_max) { + PyErr_SetString(PyExc_ValueError, "clip: must provide min and/or max"); + return NULL; + } - if (!unwrap_double(maxval_op, &maxval)) { - PyErr_SetString(PyExc_TypeError, "Expected a number"); - return NULL; - } + double minval = -INFINITY; + double maxval = INFINITY; + if (has_min && !unwrap_double(minval_op, &minval)) { + PyErr_SetString(PyExc_TypeError, "Expected a number"); + return NULL; + } + if (has_max && !unwrap_double(maxval_op, &maxval)) { + PyErr_SetString(PyExc_TypeError, "Expected a number"); + return NULL; } - if (maxval < minval) { + if (has_min && has_max && maxval < minval) { PyErr_SetString(PyExc_AssertionError, "maxval < minval"); return NULL; } @@ -2633,30 +2868,57 @@ static PyObject *Matrix_copy(PyObject *op, PyObject *Py_UNUSED(dummy)) { return (PyObject *)wrap_impl_or_free(copy); } -PyObject *Matrix_select(PyObject *op, PyObject *args) { - MatrixObject *self = (MatrixObject *)op; - matrix_impl *impl = self->impl; - - if (!impl_check_acquired(impl, true)) { - return NULL; +/// @brief Resolve one fancy-index item to an in-bounds row/column number. +/// @details Applies Python-style negative wrapping (``-1`` -> ``dim-1``) +/// then an explicit bounds check, so a list/tuple subscript can +/// never read outside the matrix. ``axis_name`` ("row"/"column") +/// and the *original* (pre-wrap) value are named in the +/// ``IndexError`` so a bad subscript is self-describing. +/// ``bool`` items are accepted as ``0``/``1`` (bool is an ``int`` +/// subclass), matching the scalar-int subscript path. +/// @return 0 with ``*out`` set on success; -1 with an exception set on +/// failure: ``TypeError`` for a non-int item, ``OverflowError`` +/// for a value outside ``Py_ssize_t`` (raised by +/// ``PyLong_AsSsize_t`` before the bounds check), or +/// ``IndexError`` for an out-of-range index. +static int resolve_gather_index(PyObject *item, size_t dim, + const char *axis_name, size_t *out) { + Py_ssize_t value = PyLong_AsSsize_t(item); + if (value == -1 && PyErr_Occurred()) { + return -1; } - PyObject *indices = NULL; - int axis = 0; - - if (!PyArg_ParseTuple(args, "O|i", &indices, &axis)) { - return NULL; + Py_ssize_t resolved = value; + if (resolved < 0) { + resolved += (Py_ssize_t)dim; } - if (axis < 0) { - axis = 2 + axis; + if (resolved < 0 || (size_t)resolved >= dim) { + PyErr_Format(PyExc_IndexError, + "%s index %zd out of range for dimension of size %zu", + axis_name, value, dim); + return -1; } - if (axis >= 2) { - PyErr_SetString(PyExc_KeyError, "Invalid axis (must be 0 or 1)"); - return NULL; - } + *out = (size_t)resolved; + return 0; +} +/// @brief Gather rows (``axis == 0``) or columns (``axis == 1``) named by a +/// sequence of indices into a fresh matrix. +/// @details Shared by ``Matrix.take`` and the subscript gather path so +/// both surfaces resolve indices identically (negative-aware and +/// bounds-checked via ``resolve_gather_index``) rather than +/// keeping two copies in sync. Row gather copies each selected +/// row with ``memcpy`` (source and destination rows are both +/// contiguous); column gather copies element-wise (inherently +/// strided). The result is wrapped as ``type`` so a subclass is +/// preserved. An empty ``indices`` sequence raises ``IndexError`` +/// and is rejected *before* any allocation — ``impl_new`` only +/// ``assert``s non-zero dimensions, so a ``0``-length axis must +/// never reach it. +static PyObject *gather_axis(matrix_impl *impl, PyObject *indices, int axis, + PyTypeObject *type) { const char *err_msg = "Indices must be specified as a list or a tuple of ints"; PyObject *fast = PySequence_Fast(indices, err_msg); @@ -2664,63 +2926,133 @@ PyObject *Matrix_select(PyObject *op, PyObject *args) { return NULL; } - Py_ssize_t size = PySequence_Fast_GET_SIZE(indices); - if (size <= 0) { + Py_ssize_t count = PySequence_Fast_GET_SIZE(fast); + if (count == 0) { Py_DECREF(fast); - Py_RETURN_NONE; + PyErr_SetString(PyExc_IndexError, "index sequence must not be empty"); + return NULL; } matrix_impl *out; if (axis == 0) { - out = impl_new((size_t)size, impl->columns); + out = impl_new((size_t)count, impl->columns); if (out == NULL) { Py_DECREF(fast); return NULL; } - double *dst = out->data; - for (Py_ssize_t i = 0; i < size; ++i) { + for (Py_ssize_t i = 0; i < count; ++i) { PyObject *item = PySequence_Fast_GET_ITEM(fast, i); - size_t r = PyLong_AsSize_t(item); - if ((Py_ssize_t)r == -1 && PyErr_Occurred()) { + size_t r; + if (resolve_gather_index(item, impl->rows, "row", &r) < 0) { Py_DECREF(fast); impl_free(out); return NULL; } - double *src = impl->row_ptrs[r]; - for (size_t i = 0; i < impl->columns; ++i, ++src, ++dst) { - *dst = *src; + memcpy(out->row_ptrs[i], impl->row_ptrs[r], + impl->columns * sizeof(double)); + } + } else { + out = impl_new(impl->rows, (size_t)count); + if (out == NULL) { + Py_DECREF(fast); + return NULL; + } + + for (Py_ssize_t i = 0; i < count; ++i) { + PyObject *item = PySequence_Fast_GET_ITEM(fast, i); + size_t c; + if (resolve_gather_index(item, impl->columns, "column", &c) < 0) { + Py_DECREF(fast); + impl_free(out); + return NULL; + } + + for (size_t r = 0; r < impl->rows; ++r) { + out->row_ptrs[r][i] = impl->row_ptrs[r][c]; } } + } - Py_DECREF(fast); - return wrap_impl_or_free(out); + Py_DECREF(fast); + + PyObject *result = wrap_matrix(type, out); + if (result == NULL) { + impl_free(out); } - out = impl_new(impl->rows, (size_t)size); - if (out == NULL) { - Py_DECREF(fast); + return result; +} + +/* Defined alongside Matrix_ass_subscript (the scatter machinery lives next + to the write path); forward-declared here so Matrix_put can share it. */ +static int scatter_axis(matrix_impl *impl, PyObject *indices, int axis, + PyObject *value_op, int accumulate); + +PyObject *Matrix_take(PyObject *op, PyObject *args, PyObject *kwds) { + MatrixObject *self = (MatrixObject *)op; + matrix_impl *impl = self->impl; + + if (!impl_check_acquired(impl, true)) { return NULL; } - size_t dst_c = 0; - for (Py_ssize_t i = 0; i < size; ++i, ++dst_c) { - PyObject *item = PySequence_Fast_GET_ITEM(fast, i); - size_t src_c = PyLong_AsSize_t(item); - if ((Py_ssize_t)src_c == -1 && PyErr_Occurred()) { - Py_DECREF(fast); - impl_free(out); - return NULL; - } + PyObject *indices = NULL; + int axis = 0; + static char *keywords[] = {"indices", "axis", NULL}; - for (size_t r = 0; r < impl->rows; ++r) { - out->row_ptrs[r][dst_c] = impl->row_ptrs[r][src_c]; - } + if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|i", keywords, &indices, + &axis)) { + return NULL; } - Py_DECREF(fast); - return wrap_impl_or_free(out); + if (axis < 0) { + axis = 2 + axis; + } + + if (axis < 0 || axis >= 2) { + PyErr_SetString(PyExc_KeyError, "Invalid axis (must be 0 or 1)"); + return NULL; + } + + return gather_axis(impl, indices, axis, Py_TYPE(self)); +} + +PyObject *Matrix_put(PyObject *op, PyObject *args, PyObject *kwds) { + MatrixObject *self = (MatrixObject *)op; + matrix_impl *impl = self->impl; + + if (!impl_check_acquired(impl, true)) { + return NULL; + } + + PyObject *indices = NULL; + PyObject *value = NULL; + int axis = 0; + int accumulate = 0; + static char *keywords[] = {"indices", "value", "axis", "accumulate", NULL}; + + if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|ip", keywords, &indices, + &value, &axis, &accumulate)) { + return NULL; + } + + if (axis < 0) { + axis = 2 + axis; + } + + if (axis < 0 || axis >= 2) { + PyErr_SetString(PyExc_KeyError, "Invalid axis (must be 0 or 1)"); + return NULL; + } + + if (scatter_axis(impl, indices, axis, value, accumulate) < 0) { + return NULL; + } + + Py_INCREF(self); + return (PyObject *)self; } static PyObject *Matrix_allclose(PyObject *cls, PyObject *args, @@ -2748,6 +3080,102 @@ static PyObject *Matrix_allclose(PyObject *cls, PyObject *args, Py_RETURN_FALSE; } +/// @brief Resolve a where() value operand to a scalar or a same-shape matrix. +/// @details ``op`` may be a Python scalar (bound into ``*scalar`` with +/// ``*mat`` left NULL), a Matrix, or a list/tuple of numbers (taken +/// as a ``1xN`` row vector). A matrix or coerced sequence must match +/// ``rows`` x ``columns`` and is returned in ``*mat`` with a fresh +/// reference the caller must release. Any other type sets a +/// ``TypeError`` and a shape mismatch sets a ``ValueError``; both +/// return -1. +static int where_resolve_operand(PyObject *op, size_t rows, size_t columns, + double *scalar, matrix_impl **mat) { + *mat = NULL; + if (unwrap_double(op, scalar)) { + return 0; + } + if (Py_TYPE(op) != LOCAL_STATE->matrix_type && !PyList_Check(op) && + !PyTuple_Check(op)) { + PyErr_SetString(PyExc_TypeError, + "where() operands must each be a Matrix, a scalar, or a " + "list/tuple of numbers"); + return -1; + } + matrix_impl *impl = unwrap_matrix(op, false); + if (impl == NULL) { + return -1; + } + if (impl->rows != rows || impl->columns != columns) { + IMPL_DECREF(impl); + PyErr_SetString(PyExc_ValueError, + "where() matrix or list/tuple operands must match the " + "mask shape"); + return -1; + } + *mat = impl; + return 0; +} + +/// @brief Select between two operands element-wise on a truthy mask. +/// @details ``where(mask, a, b)`` returns a fresh matrix taking ``a`` where the +/// corresponding ``mask`` element is non-zero and ``b`` elsewhere. +/// ``a`` and ``b`` may each be a scalar, a matrix matching the mask's +/// shape, or a list/tuple of numbers (taken as a ``1xN`` row vector +/// that must then match the mask shape); other shapes raise +/// ``ValueError``. A 1x1 matrix is treated as a matrix (not a scalar) +/// and so must match the mask shape. NaN mask elements are non-zero +/// and select ``a``. +static PyObject *Matrix_where(PyObject *cls, PyObject *args) { + PyObject *mask_op = NULL; + PyObject *a_op = NULL; + PyObject *b_op = NULL; + + if (!PyArg_ParseTuple(args, "O!OO", cls, &mask_op, &a_op, &b_op)) { + return NULL; + } + + matrix_impl *mask = ((MatrixObject *)mask_op)->impl; + if (!impl_check_acquired(mask, true)) { + return NULL; + } + const size_t rows = mask->rows; + const size_t columns = mask->columns; + + double a_scalar = 0.0; + double b_scalar = 0.0; + matrix_impl *a_mat = NULL; + matrix_impl *b_mat = NULL; + matrix_impl *out = NULL; + PyObject *result = NULL; + + if (where_resolve_operand(a_op, rows, columns, &a_scalar, &a_mat) < 0) { + goto done; + } + if (where_resolve_operand(b_op, rows, columns, &b_scalar, &b_mat) < 0) { + goto done; + } + + out = impl_new(rows, columns); + if (out == NULL) { + goto done; + } + + const double *mp = mask->data; + double *outp = out->data; + for (size_t i = 0; i < out->size; ++i, ++mp, ++outp) { + const double av = a_mat != NULL ? a_mat->data[i] : a_scalar; + const double bv = b_mat != NULL ? b_mat->data[i] : b_scalar; + *outp = (*mp != 0.0) ? av : bv; + } + + result = wrap_impl_or_free(out); + +done: + IMPL_DECREF(a_mat); + IMPL_DECREF(b_mat); + return result; +} + static int parse_dims(PyObject *shape, size_t *rows, size_t *columns) { Py_ssize_t srows; Py_ssize_t scolumns; @@ -3155,6 +3583,13 @@ static PyObject *Matrix_reduce(PyObject *op, PyObject *Py_UNUSED(dummy)) { (Py_ssize_t)impl->columns, payload); } +static PyObject *Matrix_Less_compare(PyObject *self, PyObject *other); +static PyObject *Matrix_LessEqual_compare(PyObject *self, PyObject *other); +static PyObject *Matrix_Greater_compare(PyObject *self, PyObject *other); +static PyObject *Matrix_GreaterEqual_compare(PyObject *self, PyObject *other); +static PyObject *Matrix_Equal_compare(PyObject *self, PyObject *other); +static PyObject *Matrix_NotEqual_compare(PyObject *self, PyObject *other); + static PyMethodDef Matrix_methods[] = { {"transpose", (PyCFunction)Matrix_transpose, METH_VARARGS | METH_KEYWORDS, "transpose($self, /, in_place=False)\n--\n\n" @@ -3176,6 +3611,13 @@ static PyMethodDef Matrix_methods[] = { "Axis-aware inner product: sum of element-wise products. " "Equivalent to numpy.linalg.vecdot for 1-D inputs with axis=None; " "**not** equivalent to numpy.dot."}, + {"fma", (PyCFunction)Matrix_fma, METH_VARARGS | METH_KEYWORDS, + "fma($self, b, c, /, in_place=False)\n--\n\n" + "Fused multiply-add: single-rounding ``self*b + c`` (libc fma()). " + "b and c may each be a same-shape matrix, a 1x1 matrix, a row or column " + "vector that broadcasts, or a scalar; other shapes raise ValueError. " + "The single rounding differs from ``self*b + c`` (which rounds twice) by " + "up to half a ULP, so compare results with allclose(), not ==."}, {"cross", (PyCFunction)Matrix_cross, METH_VARARGS | METH_KEYWORDS, "cross($self, other, /, axis=None)\n--\n\n" "2D (scalar z-component) or 3D cross product against another " @@ -3242,22 +3684,97 @@ static PyMethodDef Matrix_methods[] = { {"abs", (PyCFunction)Matrix_Abs_method, METH_VARARGS | METH_KEYWORDS, "abs($self, /, in_place=False)\n--\n\n" "Element-wise absolute value."}, - {"clip", Matrix_clip, METH_VARARGS, - "clip($self, min_or_maxval, /, maxval=None)\n--\n\n" - "Clip elements to a range."}, + {"sqrt", (PyCFunction)Matrix_Sqrt_method, METH_VARARGS | METH_KEYWORDS, + "sqrt($self, /, in_place=False)\n--\n\n" + "Element-wise square root. Negative elements yield NaN."}, + {"less", (PyCFunction)Matrix_Less_compare, METH_O, + "less($self, other, /)\n--\n\n" + "Element-wise ``self < other`` as a 0/1 mask matrix. ``other`` may be a " + "same-shape matrix, a scalar (including bool), a 1x1 matrix, a " + "row/column vector that broadcasts, or a list/tuple of numbers. NaN " + "comparisons yield 0. Distinct from the ``<`` operator, which returns a " + "single bool."}, + {"less_equal", (PyCFunction)Matrix_LessEqual_compare, METH_O, + "less_equal($self, other, /)\n--\n\n" + "Element-wise ``self <= other`` as a 0/1 mask matrix. ``other`` may be a " + "same-shape matrix, a scalar (including bool), a 1x1 matrix, a " + "row/column vector that broadcasts, or a list/tuple of numbers. NaN " + "comparisons yield 0. Distinct from the ``<=`` operator, which returns a " + "single bool."}, + {"greater", (PyCFunction)Matrix_Greater_compare, METH_O, + "greater($self, other, /)\n--\n\n" + "Element-wise ``self > other`` as a 0/1 mask matrix. ``other`` may be a " + "same-shape matrix, a scalar (including bool), a 1x1 matrix, a " + "row/column vector that broadcasts, or a list/tuple of numbers. NaN " + "comparisons yield 0. Distinct from the ``>`` operator, which returns a " + "single bool."}, + {"greater_equal", (PyCFunction)Matrix_GreaterEqual_compare, METH_O, + "greater_equal($self, other, /)\n--\n\n" + "Element-wise ``self >= other`` as a 0/1 mask matrix. ``other`` may be a " + "same-shape matrix, a scalar (including bool), a 1x1 matrix, a " + "row/column vector that broadcasts, or a list/tuple of numbers. NaN " + "comparisons yield 0. Distinct from the ``>=`` operator, which returns a " + "single bool."}, + {"equal", (PyCFunction)Matrix_Equal_compare, METH_O, + "equal($self, other, /)\n--\n\n" + "Element-wise ``self == other`` as a 0/1 mask matrix. ``other`` may be a " + "same-shape matrix, a scalar (including bool), a 1x1 matrix, a " + "row/column vector that broadcasts, or a list/tuple of numbers. NaN " + "comparisons yield 0. Distinct from the ``==`` operator, which returns a " + "single bool."}, + {"not_equal", (PyCFunction)Matrix_NotEqual_compare, METH_O, + "not_equal($self, other, /)\n--\n\n" + "Element-wise ``self != other`` as a 0/1 mask matrix. ``other`` may be a " + "same-shape matrix, a scalar (including bool), a 1x1 matrix, a " + "row/column vector that broadcasts, or a list/tuple of numbers. NaN " + "comparisons yield 1. Distinct from the ``!=`` operator, which returns a " + "single bool."}, + {"clip", (PyCFunction)Matrix_clip, METH_VARARGS | METH_KEYWORDS, + "clip($self, min=None, max=None)\n--\n\n" + "Clamp elements to [min, max]; either bound may be omitted.\n\n" + "The first argument is the lower bound and the second the upper\n" + "bound. Pass None (or omit) to leave a side unbounded, so\n" + "clip(min=0) clamps only below and clip(max=255) only above. Raises\n" + "ValueError if both bounds are omitted."}, {"copy", Matrix_copy, METH_NOARGS, "copy($self, /)\n--\n\nReturn a deep copy."}, {"__reduce__", Matrix_reduce, METH_NOARGS, "__reduce__($self, /)\n--\n\n" "Pickle helper: serialize the matrix to its raw double buffer."}, - {"select", Matrix_select, METH_VARARGS, - "select($self, indices, /, axis=0)\n--\n\n" - "Select rows or columns by index."}, + {"take", (PyCFunction)Matrix_take, METH_VARARGS | METH_KEYWORDS, + "take($self, indices, axis=0)\n--\n\n" + "Take rows or columns by index into a new matrix.\n\n" + "indices is a 1-D list or tuple of ints selecting whole rows\n" + "(axis=0) or columns (axis=1). Negative indices count from the\n" + "end; duplicates repeat the row or column; an out-of-range index\n" + "raises IndexError, and an empty index sequence raises IndexError."}, + {"put", (PyCFunction)Matrix_put, METH_VARARGS | METH_KEYWORDS, + "put($self, indices, value, axis=0, accumulate=False)\n--\n\n" + "Assign value into the selected rows or columns in place.\n\n" + "The write-side counterpart of take(): value may be a scalar, a\n" + "1x1 matrix, or a matrix matching the selection shape. All indices\n" + "and the value shape are validated before any write, so a rejected\n" + "call leaves the matrix unchanged. Negative indices count from the\n" + "end; an out-of-range or empty index raises IndexError. With\n" + "accumulate=False (the default) duplicate indices follow\n" + "last-write-wins; with accumulate=True the values are added into the\n" + "selection, so duplicate indices fold additively (scatter-add).\n" + "Returns self."}, {"allclose", (PyCFunction)Matrix_allclose, METH_VARARGS | METH_KEYWORDS | METH_CLASS, "allclose($type, lhs, rhs, /, rtol=1e-05, atol=1e-08, " "equal_nan=False)\n--\n\n" "Check element-wise equality within tolerance."}, + {"where", (PyCFunction)Matrix_where, METH_VARARGS | METH_CLASS, + "where($type, mask, a, b, /)\n--\n\n" + "Select element-wise from a or b on a truthy mask.\n\n" + "Returns a fresh matrix taking a where the corresponding mask\n" + "element is non-zero and b elsewhere. a and b may each be a scalar,\n" + "a matrix matching the mask's shape, or a list/tuple of numbers\n" + "(taken as a 1xN row vector that must then match the mask shape);\n" + "other shapes raise ValueError. A 1x1 matrix is treated as a matrix\n" + "(not a scalar) and so must match the mask shape. NaN mask elements\n" + "are non-zero and select a."}, {"zeros", Matrix_zeros, METH_VARARGS | METH_CLASS, "zeros($type, size, /)\n--\n\nCreate a zero-filled matrix."}, {"ones", Matrix_ones, METH_VARARGS | METH_CLASS, @@ -3497,7 +4014,7 @@ static PyGetSetDef Matrix_getset[] = { {NULL} /* Sentinel */ }; -inline enum BinaryOps swap_right(enum BinaryOps op) { +static inline enum BinaryOps swap_right(enum BinaryOps op) { switch (op) { case Subtract: return RSubtract; @@ -3505,6 +4022,18 @@ inline enum BinaryOps swap_right(enum BinaryOps op) { case Divide: return RDivide; + case Less: + return Greater; + + case Greater: + return Less; + + case LessEqual: + return GreaterEqual; + + case GreaterEqual: + return LessEqual; + default: return op; } @@ -3563,7 +4092,7 @@ static int Matrix_binary_op(PyObject *lhs_op, PyObject *rhs_op, } if (lhs->size == 1) { if (inplace) { - PyErr_SetString(PyExc_NotImplementedError, + PyErr_SetString(PyExc_ValueError, "in-place scalar broadcast would change operand shape"); goto error; } @@ -3589,12 +4118,12 @@ static int Matrix_binary_op(PyObject *lhs_op, PyObject *rhs_op, rowvec = lhs; op = swap_right(op); } else { - PyErr_SetString(PyExc_NotImplementedError, mismatch_error); + PyErr_SetString(PyExc_ValueError, mismatch_error); goto error; } if (inplace) { - PyErr_SetString(PyExc_NotImplementedError, + PyErr_SetString(PyExc_ValueError, "in-place outer broadcast would change operand shape"); goto error; } @@ -3628,7 +4157,7 @@ static int Matrix_binary_op(PyObject *lhs_op, PyObject *rhs_op, matrix = lhs; vector = rhs; } else { - PyErr_SetString(PyExc_NotImplementedError, mismatch_error); + PyErr_SetString(PyExc_ValueError, mismatch_error); goto error; } @@ -3654,7 +4183,7 @@ static int Matrix_binary_op(PyObject *lhs_op, PyObject *rhs_op, matrix = lhs; vector = rhs; } else { - PyErr_SetString(PyExc_NotImplementedError, mismatch_error); + PyErr_SetString(PyExc_ValueError, mismatch_error); goto error; } @@ -3703,6 +4232,20 @@ static int Matrix_binary_op(PyObject *lhs_op, PyObject *rhs_op, return out; \ } +/* MATRIX_COMPARE_METHOD stamps a named element-wise comparison method that + reuses Matrix_binary_op for the full broadcast routing (same-shape, scalar, + 1x1, row-vector, column-vector). Always invoked as self.(other), so + self is the left operand and other the right. Returns a fresh 0/1 mask + matrix; the comparison kernels compile branch-free to cmppd + andpd. */ +#define MATRIX_COMPARE_METHOD(ENUM) \ + static PyObject *Matrix_##ENUM##_compare(PyObject *self, PyObject *other) { \ + PyObject *out = NULL; \ + if (Matrix_binary_op(self, other, &out, ENUM, false) < 0) { \ + return NULL; \ + } \ + return out; \ + } + /* MATRIX_UNARY_OP stamps a Python number-protocol slot wrapper that calls the per-op kernel impl__ewise directly. Only stamped for ops with a number-protocol slot (Py_nb_absolute, Py_nb_negative). */ @@ -3733,6 +4276,168 @@ MATRIX_INPLACE_BINARY_OP(Divide) MATRIX_UNARY_OP(Abs, abs) MATRIX_UNARY_OP(Negate, negate) +MATRIX_COMPARE_METHOD(Less) +MATRIX_COMPARE_METHOD(LessEqual) +MATRIX_COMPARE_METHOD(Greater) +MATRIX_COMPARE_METHOD(GreaterEqual) +MATRIX_COMPARE_METHOD(Equal) +MATRIX_COMPARE_METHOD(NotEqual) + +/// @brief Lexicographic three-way compare of two equal-length buffers. +/// @details Walks both buffers in row-major order and returns at the first +/// element where a strict ``<`` or ``>`` holds: ``-1`` if ``lhs`` is +/// smaller there, ``+1`` if larger. Elements that are equal — or +/// where neither ordering holds, as with NaN — are skipped, so a +/// run of NaNs does not decide the result. Returns ``0`` when no +/// element decides (the buffers compare equal element-wise). +static int impl_lexcompare(const double *lhs, const double *rhs, size_t n) { + for (size_t i = 0; i < n; ++i) { + if (lhs[i] < rhs[i]) { + return -1; + } + if (lhs[i] > rhs[i]) { + return 1; + } + } + return 0; +} + +/// @brief Map a three-way compare result and a Py_RichCompare opid to a bool. +static PyObject *richcompare_from_sign(int sign, int opid) { + int result; + switch (opid) { + case Py_LT: + result = sign < 0; + break; + case Py_LE: + result = sign <= 0; + break; + case Py_GT: + result = sign > 0; + break; + case Py_GE: + result = sign >= 0; + break; + case Py_EQ: + result = sign == 0; + break; + case Py_NE: + result = sign != 0; + break; + default: + Py_RETURN_NOTIMPLEMENTED; + } + return PyBool_FromLong(result); +} + +/// @brief Lexicographic compare of two matrices, total for ``==`` / ``!=``. +/// @details A shape mismatch makes ``==`` ``False`` and ``!=`` ``True``; the +/// ordering operators raise ``ValueError`` instead. Equal-shape +/// buffers compare lexicographically in row-major order. +static PyObject *richcompare_matrix(const matrix_impl *lhs, + const matrix_impl *rhs, int opid) { + if (lhs->rows != rhs->rows || lhs->columns != rhs->columns) { + if (opid == Py_EQ) { + Py_RETURN_FALSE; + } + if (opid == Py_NE) { + Py_RETURN_TRUE; + } + PyErr_SetString(PyExc_ValueError, + "ordering comparison requires a matrix, list, or tuple " + "operand of the same shape (or a scalar operand)"); + return NULL; + } + int sign = impl_lexcompare(lhs->data, rhs->data, lhs->size); + return richcompare_from_sign(sign, opid); +} + +/// @brief Lexicographic ``bool`` comparison operators for Matrix. +/// @details Returns a single Python ``bool``, comparing element by element in +/// row-major order like a list or tuple: the first element where a +/// strict ordering holds decides, and all-equal compares equal. The +/// right operand may be a same-shape matrix, a Python scalar (which +/// broadcasts to ``self``'s shape), or a list/tuple of numbers (taken +/// as a ``1xN`` row vector that must then match ``self``'s shape). +/// ``==`` / ``!=`` are total: a shape mismatch makes them +/// ``False`` / ``True``, an operand that is not a matrix, scalar, or +/// list/tuple yields ``NotImplemented`` so Python falls back to +/// identity, and a list/tuple that cannot be coerced (empty, or +/// containing a non-number) likewise makes ``==`` ``False`` / +/// ``!=`` ``True`` instead of raising — so ``matrix in some_list`` +/// still works. They are not total over ownership: comparing a matrix +/// not owned by the current interpreter raises ``RuntimeError``. +/// Defining value equality makes Matrix +/// unhashable (CPython disables the inherited ``__hash__``), which is +/// correct for a mutable type. The ordering operators +/// (``<`` ``<=`` ``>`` ``>=``) raise ``ValueError`` on a shape +/// mismatch (and propagate the coercion error — ``ValueError`` for an +/// empty list, ``TypeError`` for a non-number element). NaN never +/// decides an ordering (neither ``<`` nor ``>`` +/// holds, so the element is skipped); an all-NaN matrix therefore +/// compares ``==`` equal to itself — unlike the ``equal`` mask, where +/// ``NaN == x`` is ``0.0``. A 1x1 matrix is NOT treated as a scalar +/// here (unlike the arithmetic and mask paths); it only compares +/// against another 1x1. +static PyObject *Matrix_richcompare(PyObject *self_op, PyObject *other_op, + int opid) { + MatrixObject *self = (MatrixObject *)self_op; + matrix_impl *lhs = self->impl; + if (!impl_check_acquired(lhs, true)) { + return NULL; + } + + double scalar; + if (unwrap_double(other_op, &scalar)) { + int sign = 0; + for (size_t i = 0; i < lhs->size; ++i) { + if (lhs->data[i] < scalar) { + sign = -1; + break; + } + if (lhs->data[i] > scalar) { + sign = 1; + break; + } + } + return richcompare_from_sign(sign, opid); + } + + if (Py_TYPE(other_op) == LOCAL_STATE->matrix_type) { + matrix_impl *rhs = ((MatrixObject *)other_op)->impl; + if (!impl_check_acquired(rhs, true)) { + return NULL; + } + return richcompare_matrix(lhs, rhs, opid); + } + + if (PyList_Check(other_op) || PyTuple_Check(other_op)) { + matrix_impl *coerced = impl_new_from_sequence(other_op, false); + if (coerced == NULL) { + // == / != stay total over coercion failure: an empty or non-numeric + // list/tuple counts as "not equal" rather than raising, so + // ``matrix in some_list`` works. Only ValueError/TypeError are swallowed; + // a MemoryError still propagates so an alloc failure is never read as + // equal. + if ((opid == Py_EQ || opid == Py_NE) && + (PyErr_ExceptionMatches(PyExc_ValueError) || + PyErr_ExceptionMatches(PyExc_TypeError))) { + PyErr_Clear(); + if (opid == Py_EQ) { + Py_RETURN_FALSE; + } + Py_RETURN_TRUE; + } + return NULL; + } + PyObject *result = richcompare_matrix(lhs, coerced, opid); + impl_free(coerced); + return result; + } + + Py_RETURN_NOTIMPLEMENTED; +} + static PyObject *Matrix_matmul(PyObject *lhs_op, PyObject *rhs_op) { matrix_impl *lhs = NULL; matrix_impl *rhs = NULL; @@ -3750,7 +4455,7 @@ static PyObject *Matrix_matmul(PyObject *lhs_op, PyObject *rhs_op) { } if (lhs->columns != rhs->rows) { - PyErr_SetString(PyExc_NotImplementedError, "M0xN0 @ M1xN1 N0 != M1"); + PyErr_SetString(PyExc_ValueError, "M0xN0 @ M1xN1 N0 != M1"); goto error; } @@ -3821,6 +4526,62 @@ static int read_key_ranges(PyObject *key, range *rows, range *columns, return 0; } +/// @brief True iff ``obj`` is the bare ``slice(None, None, None)`` (``:``) +/// @details Only the bare ``:`` qualifies as a full selector alongside a list +/// key; an explicit range such as ``0:R`` does not. +static int is_full_slice(PyObject *obj) { + if (!PySlice_Check(obj)) { + return 0; + } + + PySliceObject *slice = (PySliceObject *)obj; + return slice->start == Py_None && slice->stop == Py_None && + slice->step == Py_None; +} + +/// @brief How a subscript key maps onto list-based fancy indexing. +typedef enum { + FANCY_KEY_NONE = 0, ///< Not a fancy key; fall through to range parsing. + FANCY_KEY_AXIS0, ///< A list selects whole rows. + FANCY_KEY_AXIS1, ///< A list selects whole columns. + FANCY_KEY_UNSUPPORTED ///< A list-bearing key in an unsupported shape. +} fancy_key_kind; + +/// @brief Classify a subscript key for list-based fancy indexing +/// @details Shared by the read and write paths so both accept the same +/// shapes: ``m[[rows]]``, ``m[[rows], :]``, ``m[:, [cols]]``. The +/// int/slice hot path returns ``FANCY_KEY_NONE``. On +/// ``FANCY_KEY_UNSUPPORTED`` no exception is set — the caller raises. +/// @param key The subscript key. +/// @param list_out On an axis match, set to the borrowed index-list object. +/// @return The fancy-key classification. +static fancy_key_kind classify_fancy_key(PyObject *key, PyObject **list_out) { + if (PyList_Check(key)) { + *list_out = key; + return FANCY_KEY_AXIS0; + } + + if (PyTuple_Check(key) && PyTuple_GET_SIZE(key) == 2) { + PyObject *k0 = PyTuple_GET_ITEM(key, 0); + PyObject *k1 = PyTuple_GET_ITEM(key, 1); + int list0 = PyList_Check(k0); + int list1 = PyList_Check(k1); + if (list0 || list1) { + if (list0 && !list1 && is_full_slice(k1)) { + *list_out = k0; + return FANCY_KEY_AXIS0; + } + if (list1 && !list0 && is_full_slice(k0)) { + *list_out = k1; + return FANCY_KEY_AXIS1; + } + return FANCY_KEY_UNSUPPORTED; + } + } + + return FANCY_KEY_NONE; +} + PyObject *Matrix_subscript(PyObject *op, PyObject *key) { MatrixObject *self = (MatrixObject *)op; matrix_impl *impl = self->impl; @@ -3829,6 +4590,23 @@ PyObject *Matrix_subscript(PyObject *op, PyObject *key) { return NULL; } + // Dispatch fancy indexing before read_key_ranges; non-list keys classify + // as FANCY_KEY_NONE and fall through unchanged. + PyObject *index_list; + switch (classify_fancy_key(key, &index_list)) { + case FANCY_KEY_AXIS0: + return gather_axis(impl, index_list, 0, Py_TYPE(op)); + case FANCY_KEY_AXIS1: + return gather_axis(impl, index_list, 1, Py_TYPE(op)); + case FANCY_KEY_UNSUPPORTED: + PyErr_SetString(PyExc_IndexError, + "fancy indexing supports only m[[rows]], m[[rows], :], or " + "m[:, [cols]]; use take() for other selections"); + return NULL; + case FANCY_KEY_NONE: + break; + } + range rows; range columns; if (read_key_ranges(key, &rows, &columns, impl) < 0) { @@ -3885,6 +4663,275 @@ static PyObject *Matrix_item(PyObject *op, Py_ssize_t index) { static PyObject *Matrix_iter(PyObject *op) { return PySeqIter_New(op); } +/// @brief Resolve an entire index list into a validated ``size_t`` array. +/// @details The scatter (write) counterpart of ``gather_axis``'s per-element +/// resolution. Every index is validated (negative-aware, +/// bounds-checked) *before* the caller writes any element, so the +/// write phase is infallible — the property that licenses the +/// ``memcpy`` row kernel with no mid-loop error exit. Indices are +/// written into ``stack`` when ``count <= stack_size`` (zero heap +/// allocation for the common small scatter) and into a fresh heap +/// buffer otherwise; on success ``*out_idx`` points at whichever +/// buffer was used, so the caller frees it only when it differs +/// from ``stack``. An empty sequence raises ``IndexError`` +/// (matching gather). On any failure the heap buffer, if one was +/// allocated, is freed here before returning. +/// @param indices The list/sequence of index objects. +/// @param dim Size of the axis being indexed, for bounds checking. +/// @param axis_name ``"row"`` or ``"column"`` for error messages. +/// @param stack Caller-provided fixed buffer used when small enough. +/// @param stack_size Number of slots in ``stack``. +/// @param out_idx Receives the resolved index array (``stack`` or heap). +/// @param out_count Receives the number of resolved indices. +/// @return 0 on success; -1 with an exception set on failure. +static int resolve_index_list(PyObject *indices, size_t dim, + const char *axis_name, size_t *stack, + size_t stack_size, size_t **out_idx, + size_t *out_count) { + const char *err_msg = + "Indices must be specified as a list or a tuple of ints"; + PyObject *fast = PySequence_Fast(indices, err_msg); + if (fast == NULL) { + return -1; + } + + Py_ssize_t count = PySequence_Fast_GET_SIZE(fast); + if (count == 0) { + Py_DECREF(fast); + PyErr_SetString(PyExc_IndexError, "index sequence must not be empty"); + return -1; + } + + size_t *idx = stack; + if ((size_t)count > stack_size) { + idx = PyMem_Malloc((size_t)count * sizeof(size_t)); + if (idx == NULL) { + Py_DECREF(fast); + PyErr_NoMemory(); + return -1; + } + } + + for (Py_ssize_t i = 0; i < count; ++i) { + PyObject *item = PySequence_Fast_GET_ITEM(fast, i); + if (resolve_gather_index(item, dim, axis_name, &idx[i]) < 0) { + if (idx != stack) { + PyMem_Free(idx); + } + Py_DECREF(fast); + return -1; + } + } + + Py_DECREF(fast); + *out_idx = idx; + *out_count = (size_t)count; + return 0; +} + +/// @brief Classify a scatter RHS as a scalar or an exact-shape matrix. +/// @details Mirrors ``classify_fma_operand``: a real number or a ``1x1`` +/// matrix is a broadcast scalar (the in-house scalar rule), and a +/// matrix matching the selection shape exactly is INCREF'd into +/// ``*full``. Any other matrix shape raises ``ValueError`` naming +/// the offered shape and the selection shape. ``want_rows`` / +/// ``want_cols`` are the shape the selection writes (``count`` by +/// the receiver's column count for a row scatter, the receiver's +/// row count by ``count`` for a column scatter). +/// @param value_op The assigned value object. +/// @param want_rows Expected RHS row count for a full-matrix RHS. +/// @param want_cols Expected RHS column count for a full-matrix RHS. +/// @param scalar Receives the broadcast value when the RHS is a scalar. +/// @param full Receives an INCREF'd impl when the RHS is a full matrix, +/// otherwise ``NULL``. +/// @return 0 on success; -1 with an exception set on failure. +static int classify_scatter_rhs(PyObject *value_op, size_t want_rows, + size_t want_cols, double *scalar, + matrix_impl **full) { + *scalar = 0.0; + *full = NULL; + + if (Py_TYPE(value_op) == LOCAL_STATE->matrix_type) { + matrix_impl *impl = ((MatrixObject *)value_op)->impl; + if (!impl_check_acquired(impl, true)) { + return -1; + } + if (impl->size == 1) { + *scalar = impl->data[0]; + return 0; + } + if (impl->rows == want_rows && impl->columns == want_cols) { + IMPL_INCREF(impl); + *full = impl; + return 0; + } + PyErr_Format(PyExc_ValueError, + "cannot assign %zux%zu matrix to a selection of shape %zux%zu", + impl->rows, impl->columns, want_rows, want_cols); + return -1; + } + + if (unwrap_double(value_op, scalar)) { + return 0; + } + if (PyErr_Occurred()) { + return -1; + } + + PyErr_SetString(PyExc_TypeError, "value must be a Matrix or a real number"); + return -1; +} + +/// @brief Scatter a same-shape matrix into the selected rows +/// @details ``value`` is ``count x columns``. When ``accumulate`` each +/// source row is added into its destination; otherwise it overwrites. +static void impl_scatter_rows(matrix_impl *matrix, const size_t *idx, + size_t count, matrix_impl *value, + int accumulate) { + for (size_t i = 0; i < count; ++i) { + double *dst = matrix->row_ptrs[idx[i]]; + const double *src = value->row_ptrs[i]; + if (accumulate) { + for (size_t c = 0; c < matrix->columns; ++c) { + dst[c] += src[c]; + } + } else { + memcpy(dst, src, matrix->columns * sizeof(double)); + } + } +} + +/// @brief Broadcast a scalar into the selected rows +static void impl_scatter_rows_scalar(matrix_impl *matrix, const size_t *idx, + size_t count, double value, + int accumulate) { + for (size_t i = 0; i < count; ++i) { + double *dst = matrix->row_ptrs[idx[i]]; + for (size_t c = 0; c < matrix->columns; ++c) { + if (accumulate) { + dst[c] += value; + } else { + dst[c] = value; + } + } + } +} + +/// @brief Scatter a same-shape matrix into the selected columns (strided) +/// @details ``value`` is ``rows x count``. When ``accumulate`` each source +/// element is added into its destination; otherwise it overwrites. +static void impl_scatter_cols(matrix_impl *matrix, const size_t *idx, + size_t count, matrix_impl *value, + int accumulate) { + for (size_t r = 0; r < matrix->rows; ++r) { + double *dst = matrix->row_ptrs[r]; + const double *src = value->row_ptrs[r]; + for (size_t i = 0; i < count; ++i) { + if (accumulate) { + dst[idx[i]] += src[i]; + } else { + dst[idx[i]] = src[i]; + } + } + } +} + +/// @brief Broadcast a scalar into the selected columns (strided) +static void impl_scatter_cols_scalar(matrix_impl *matrix, const size_t *idx, + size_t count, double value, + int accumulate) { + for (size_t r = 0; r < matrix->rows; ++r) { + double *dst = matrix->row_ptrs[r]; + for (size_t i = 0; i < count; ++i) { + if (accumulate) { + dst[idx[i]] += value; + } else { + dst[idx[i]] = value; + } + } + } +} + +/// @brief Three-phase scatter-store for a list-keyed assignment +/// @details Resolve all indices, then classify/shape-check the RHS, then +/// write — so a rejected scatter leaves the receiver unmodified. +/// When ``accumulate`` the RHS is added in (else it overwrites). A +/// self-aliased matrix RHS is snapshotted before the write. +/// @pre Callers must ``impl_check_acquired(impl)`` before calling; the RHS is +/// gated internally by ``classify_scatter_rhs``. +/// @param impl The receiver matrix impl (written in place). +/// @param indices The list of row (``axis == 0``) or column (``axis == 1``) +/// indices to assign. +/// @param axis 0 to scatter rows, 1 to scatter columns. +/// @param value_op The assigned scalar or matrix. +/// @param accumulate When non-zero, add into the selection instead of +/// overwriting. +/// @return 0 on success; -1 with an exception set on failure. +static int scatter_axis(matrix_impl *impl, PyObject *indices, int axis, + PyObject *value_op, int accumulate) { + size_t stack[64]; + size_t *idx; + size_t count; + size_t dim = axis == 0 ? impl->rows : impl->columns; + const char *axis_name = axis == 0 ? "row" : "column"; + if (resolve_index_list(indices, dim, axis_name, stack, + sizeof(stack) / sizeof(stack[0]), &idx, &count) < 0) { + return -1; + } + + size_t want_rows = axis == 0 ? count : impl->rows; + size_t want_cols = axis == 0 ? impl->columns : count; + double scalar; + matrix_impl *value; + if (classify_scatter_rhs(value_op, want_rows, want_cols, &scalar, &value) < + 0) { + if (idx != stack) { + PyMem_Free(idx); + } + return -1; + } + + if (value == NULL) { + if (axis == 0) { + impl_scatter_rows_scalar(impl, idx, count, scalar, accumulate); + } else { + impl_scatter_cols_scalar(impl, idx, count, scalar, accumulate); + } + } else { + // Snapshot a self-aliased RHS first: an in-place write would otherwise + // clobber cells still to be read, corrupting a permutation and aliasing + // memcpy on a fixed point. + matrix_impl *src = value; + matrix_impl *snapshot = NULL; + if (value == impl) { + snapshot = impl_new(value->rows, value->columns); + if (snapshot == NULL) { + IMPL_DECREF(value); + if (idx != stack) { + PyMem_Free(idx); + } + return -1; + } + memcpy(snapshot->data, value->data, value->size * sizeof(double)); + src = snapshot; + } + if (axis == 0) { + impl_scatter_rows(impl, idx, count, src, accumulate); + } else { + impl_scatter_cols(impl, idx, count, src, accumulate); + } + if (snapshot != NULL) { + impl_free(snapshot); + } + IMPL_DECREF(value); + } + + if (idx != stack) { + PyMem_Free(idx); + } + return 0; +} + int Matrix_ass_subscript(PyObject *op, PyObject *key, PyObject *value_op) { MatrixObject *self = (MatrixObject *)op; matrix_impl *impl = self->impl; @@ -3893,6 +4940,24 @@ int Matrix_ass_subscript(PyObject *op, PyObject *key, PyObject *value_op) { return -1; } + // Mirror the read prologue: fancy assignment shares classify_fancy_key, so + // a FANCY_KEY_NONE key falls through to the plain assignment path below. + PyObject *index_list; + switch (classify_fancy_key(key, &index_list)) { + case FANCY_KEY_AXIS0: + return scatter_axis(impl, index_list, 0, value_op, 0); + case FANCY_KEY_AXIS1: + return scatter_axis(impl, index_list, 1, value_op, 0); + case FANCY_KEY_UNSUPPORTED: + PyErr_SetString( + PyExc_IndexError, + "fancy assignment supports only m[[rows]] = v, m[[rows], :] = v, " + "or m[:, [cols]] = v; use put() for other selections"); + return -1; + case FANCY_KEY_NONE: + break; + } + range rows; range columns; if (read_key_ranges(key, &rows, &columns, impl) < 0) { @@ -4089,6 +5154,7 @@ static PyType_Slot Matrix_slots[] = { {Py_tp_str, Matrix_str}, {Py_tp_repr, Matrix_repr}, {Py_tp_iter, Matrix_iter}, + {Py_tp_richcompare, Matrix_richcompare}, {Py_nb_add, Matrix_Add_op}, {Py_nb_inplace_add, Matrix_inplace_Add_op}, {Py_nb_subtract, Matrix_Subtract_op}, diff --git a/templates/c_abi_consumer/pyproject.toml b/templates/c_abi_consumer/pyproject.toml index 6cbdaa3..626d968 100644 --- a/templates/c_abi_consumer/pyproject.toml +++ b/templates/c_abi_consumer/pyproject.toml @@ -4,11 +4,11 @@ # so the build resolves headers against the bocpy install actually # being tested. See README.md. # -# The compatible-release bound (~=0.11) keeps the template aligned with +# The compatible-release bound (~=0.12) keeps the template aligned with # the public C ABI it was authored against; bump it in lock-step with # ``[project].version`` in the root ``pyproject.toml`` (see the # ``finalize-pr`` skill). -requires = ["setuptools", "wheel", "bocpy~=0.11"] +requires = ["setuptools", "wheel", "bocpy~=0.12"] build-backend = "setuptools.build_meta" [project] @@ -16,4 +16,4 @@ name = "bocpy-c-abi-consumer" version = "0.0.0" description = "Smoke test and canonical downstream template for the bocpy public C ABI." requires-python = ">=3.10" -dependencies = ["bocpy~=0.11"] +dependencies = ["bocpy~=0.12"] diff --git a/test/test_matrix.py b/test/test_matrix.py index 14aca94..64d941c 100644 --- a/test/test_matrix.py +++ b/test/test_matrix.py @@ -1,6 +1,7 @@ """Tests for the bocpy Matrix class using fuzzed inputs across multiple sizes.""" import copy +from fractions import Fraction import math import pickle import random @@ -14,6 +15,21 @@ QUIESCE_TIMEOUT = 5 +def ref_fma(a, b, c): + """Correctly-rounded ``a * b + c`` reference, matching C99 ``fma``. + + ``math.fma`` is only available on CPython 3.13+, but the suite must run + on the supported 3.10-3.12 legs too. Exact rational arithmetic followed + by a single round-to-nearest (``float(Fraction)`` rounds half-to-even, + the same rule as IEEE-754) reproduces the single-rounding ``fma`` result + bit-for-bit for finite operands whose exact ``a * b + c`` lies within + double's range. (An exact result above ``DBL_MAX`` raises ``OverflowError`` + here where C99 ``fma`` returns an infinity; the test inputs stay well + inside range, so this boundary never bites.) + """ + return float(Fraction(a) * Fraction(b) + Fraction(c)) + + MATRIX_SIZES = [ (1, 1), (1, 5), @@ -395,7 +411,7 @@ def test_inplace_scalar_matrix_op_matrix_rejected(self): """``1x1 op= MxN`` would change the operand shape and is rejected.""" a = Matrix(1, 1, 5.0) b = Matrix(2, 3, [1, 2, 3, 4, 5, 6]) - with pytest.raises(NotImplementedError, + with pytest.raises(ValueError, match="in-place scalar broadcast"): a += b @@ -640,7 +656,7 @@ def test_matches_magnitude_squared(self, mat, shape): def test_invalid_axis_raises(self, mat): """Out-of-range axis surfaces the same error as `magnitude`.""" - with pytest.raises(NotImplementedError, match="axis must be -2, -1, 0, or 1"): + with pytest.raises(ValueError, match="axis must be -2, -1, 0, or 1"): mat.magnitude_squared(2) @@ -766,21 +782,21 @@ def test_vector_length_mismatch_raises(self): """Mismatched vector lengths surface the dimension-mismatch error.""" a = Matrix(1, 2, [1.0, 2.0]) b = Matrix(1, 3, [3.0, 4.0, 5.0]) - with pytest.raises(NotImplementedError, match=r"vecdot: lhs \d+x\d+ incompatible with rhs \d+x\d+"): + with pytest.raises(ValueError, match=r"vecdot: lhs \d+x\d+ incompatible with rhs \d+x\d+"): a.vecdot(b) def test_incompatible_matrix_shapes_raises(self): """Incompatible matrix shapes (no broadcast match) surface the same error.""" a = Matrix(2, 3, [1.0] * 6) b = Matrix(4, 5, [1.0] * 20) - with pytest.raises(NotImplementedError, match=r"vecdot: lhs \d+x\d+ incompatible with rhs \d+x\d+"): + with pytest.raises(ValueError, match=r"vecdot: lhs \d+x\d+ incompatible with rhs \d+x\d+"): a.vecdot(b) def test_invalid_axis_raises(self): """Out-of-range axis surfaces the same error as `magnitude`.""" a = Matrix(1, 3, [1.0, 2.0, 3.0]) b = Matrix(1, 3, [4.0, 5.0, 6.0]) - with pytest.raises(NotImplementedError, match="axis must be -2, -1, 0, or 1"): + with pytest.raises(ValueError, match="axis must be -2, -1, 0, or 1"): a.vecdot(b, 2) def test_axis_wrong_type_raises(self): @@ -880,6 +896,217 @@ def test_vecdot_keyword_axis_none_matches_default(self): assert a.vecdot(b, axis=None) == pytest.approx(a.vecdot(b)) +class TestFma: + """Tests for `fma(b, c, /, in_place=False)` single-rounding multiply-add.""" + + def test_exact_small(self): + """A hand-computed 2x2 matches a single-rounding fma cell-by-cell.""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + b = Matrix(2, 2, [10.0, 20.0, 30.0, 40.0]) + c = Matrix(2, 2, [0.5, 0.5, 0.5, 0.5]) + out = a.fma(b, c) + for i in range(2): + for j in range(2): + assert out[i, j] == ref_fma(a[i, j], b[i, j], c[i, j]) + + def test_parity_same_shape(self, shape, rng): + """fma(b, c) is close to a*b + c for same-shape matrix operands.""" + rows, cols = shape + n = rows * cols + a = Matrix(rows, cols, [rng.uniform(-100, 100) for _ in range(n)]) + b = Matrix(rows, cols, [rng.uniform(-100, 100) for _ in range(n)]) + c = Matrix(rows, cols, [rng.uniform(-100, 100) for _ in range(n)]) + # One rounding vs two: compare with allclose, never ==. + assert Matrix.allclose(a.fma(b, c), a * b + c) + + def test_exact_matches_math_fma_fuzz(self, shape, rng): + """Every cell of fma(b, c) equals a single-rounding reference fma.""" + rows, cols = shape + n = rows * cols + av = [rng.uniform(-100, 100) for _ in range(n)] + bv = [rng.uniform(-100, 100) for _ in range(n)] + cv = [rng.uniform(-100, 100) for _ in range(n)] + a = Matrix(rows, cols, av) + b = Matrix(rows, cols, bv) + c = Matrix(rows, cols, cv) + out = a.fma(b, c) + for i in range(rows): + for j in range(cols): + k = i * cols + j + assert out[i, j] == ref_fma(av[k], bv[k], cv[k]) + + def test_scalar_both(self): + """Scalar b and c broadcast across every cell.""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + out = a.fma(2.0, 1.0) + for i in range(2): + for j in range(2): + assert out[i, j] == ref_fma(a[i, j], 2.0, 1.0) + + def test_scalar_1x1_is_scalar(self): + """A 1x1 matrix operand broadcasts like a scalar.""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + b = Matrix(1, 1, [2.0]) + c = Matrix(1, 1, [1.0]) + out = a.fma(b, c) + for i in range(2): + for j in range(2): + assert out[i, j] == ref_fma(a[i, j], 2.0, 1.0) + + def test_headline_mx_plus_b(self): + """The m*x + b headline form: full multiplier, scalar addend.""" + m = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + x = Matrix(2, 3, [2.0, 2.0, 2.0, 3.0, 3.0, 3.0]) + out = m.fma(x, 1.0) + assert Matrix.allclose(out, m * x + 1.0) + + def test_mixed_scalar_b_matrix_c(self): + """A scalar b with a matrix c routes through the one kernel.""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + c = Matrix(2, 2, [10.0, 20.0, 30.0, 40.0]) + out = a.fma(2.0, c) + for i in range(2): + for j in range(2): + assert out[i, j] == ref_fma(a[i, j], 2.0, c[i, j]) + + def test_mixed_matrix_b_scalar_c(self): + """A matrix b with a scalar c routes through the one kernel.""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + b = Matrix(2, 2, [10.0, 20.0, 30.0, 40.0]) + out = a.fma(b, 3.0) + for i in range(2): + for j in range(2): + assert out[i, j] == ref_fma(a[i, j], b[i, j], 3.0) + + def test_in_place_returns_self(self): + """in_place=True writes into self and returns it; b/c unchanged.""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + b = Matrix(2, 2, [10.0, 20.0, 30.0, 40.0]) + c = Matrix(2, 2, [0.5, 0.5, 0.5, 0.5]) + b_before = [b[i, j] for i in range(2) for j in range(2)] + c_before = [c[i, j] for i in range(2) for j in range(2)] + expected = [ref_fma(a[i, j], b[i, j], c[i, j]) + for i in range(2) for j in range(2)] + out = a.fma(b, c, in_place=True) + assert out is a + assert [a[i, j] for i in range(2) for j in range(2)] == expected + assert [b[i, j] for i in range(2) for j in range(2)] == b_before + assert [c[i, j] for i in range(2) for j in range(2)] == c_before + + def test_in_place_keyword_matches_copy(self, shape, rng): + """in_place result equals the copy result, leaving operands intact.""" + rows, cols = shape + n = rows * cols + av = [rng.uniform(-50, 50) for _ in range(n)] + a_copy = Matrix(rows, cols, av) + a_ip = Matrix(rows, cols, av) + b = Matrix(rows, cols, [rng.uniform(-50, 50) for _ in range(n)]) + c = Matrix(rows, cols, [rng.uniform(-50, 50) for _ in range(n)]) + copied = a_copy.fma(b, c) + a_ip.fma(b, c, in_place=True) + assert Matrix.allclose(a_ip, copied) + + def test_rejects_broadcast_b(self): + """A b matrix of incompatible shape raises ValueError.""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + bad = Matrix(3, 3, [float(i) for i in range(9)]) + with pytest.raises(ValueError, + match=r"fma: b shape 3x3 incompatible with self 2x2"): + a.fma(bad, 1.0) + + def test_rejects_broadcast_c(self): + """A c matrix of incompatible shape raises ValueError.""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + bad = Matrix(1, 3, [1.0, 2.0, 3.0]) + with pytest.raises(ValueError, + match=r"fma: c shape 1x3 incompatible with self 2x2"): + a.fma(2.0, bad) + + def test_rejects_broadcast_leaves_self_unmodified(self): + """A rejected operand allocates nothing and leaves self untouched.""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + before = [a[i, j] for i in range(2) for j in range(2)] + bad = Matrix(3, 3, [float(i) for i in range(9)]) + with pytest.raises(ValueError): + a.fma(2.0, bad, in_place=True) + assert [a[i, j] for i in range(2) for j in range(2)] == before + + def test_rejects_bad_type_b(self): + """A non-numeric, non-matrix b raises TypeError.""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + with pytest.raises(TypeError, + match=r"fma: b must be a Matrix or a real number"): + a.fma("x", 1.0) + + def test_rejects_bad_type_c(self): + """A non-numeric, non-matrix c raises TypeError.""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + with pytest.raises(TypeError, + match=r"fma: c must be a Matrix or a real number"): + a.fma(1.0, object()) + + def test_operand_aliases_self(self, shape, rng): + """fma(a, c), fma(b, a) and fma(a, a) read each cell only once.""" + rows, cols = shape + n = rows * cols + av = [rng.uniform(-20, 20) for _ in range(n)] + cv = [rng.uniform(-20, 20) for _ in range(n)] + a1 = Matrix(rows, cols, av) + c1 = Matrix(rows, cols, cv) + # b aliases self. + assert Matrix.allclose(a1.fma(a1, c1), a1 * a1 + c1) + a2 = Matrix(rows, cols, av) + # c aliases self. + assert Matrix.allclose(a2.fma(c1, a2), c1 * a2 + a2) + a3 = Matrix(rows, cols, av) + # both operands alias self, in place. + ref = Matrix(rows, cols, av) + out = a3.fma(a3, a3, in_place=True) + assert out is a3 + assert Matrix.allclose(a3, ref * ref + ref) + + def test_keyword_in_place_false_is_copy(self): + """Explicit in_place=False returns a distinct matrix.""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + before = [a[i, j] for i in range(2) for j in range(2)] + out = a.fma(2.0, 1.0, in_place=False) + assert out is not a + assert [a[i, j] for i in range(2) for j in range(2)] == before + + def test_boc_roundtrip(self): + """fma runs inside a @when behavior over a Cown[Matrix].""" + a = Cown(Matrix(2, 2, [1.0, 2.0, 3.0, 4.0])) + b = Matrix(2, 2, [10.0, 20.0, 30.0, 40.0]) + c = Matrix(2, 2, [0.5, 1.5, 2.5, 3.5]) + + @when(a) + def result(a, b=b, c=c): # noqa: D401 — short behavior + """Compute fma inside a behavior and return the result.""" + return a.value.fma(b, c) + + wait() + assert result.exception is False + expected = Matrix(2, 2, [ref_fma(v, w, x) for v, w, x in zip( + [1.0, 2.0, 3.0, 4.0], [10.0, 20.0, 30.0, 40.0], + [0.5, 1.5, 2.5, 3.5])]) + assert Matrix.allclose(result.value, expected) + + def test_not_acquired_raises(self): + """fma on a cown-resident (unacquired) matrix raises RuntimeError.""" + m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + Cown(m) + with pytest.raises(RuntimeError): + m.fma(2.0, 1.0) + + def test_operand_not_acquired_raises(self): + """An unacquired matrix operand raises RuntimeError.""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + b = Matrix(2, 2, [10.0, 20.0, 30.0, 40.0]) + Cown(b) + with pytest.raises(RuntimeError): + a.fma(b, 1.0) + + class TestCross: """Tests for the 2D / 3D ``cross`` method.""" @@ -1243,7 +1470,7 @@ def test_reverse_broadcast_vector_self_batch_other_raises(self): """Cross is anticommutative; reverse broadcast (vec.cross(batch)) is rejected.""" a = Matrix(1, 3, [1.0, 2.0, 3.0]) b = Matrix(5, 3, [float(i) for i in range(15)]) - with pytest.raises(NotImplementedError, + with pytest.raises(ValueError, match=r"cross: .* incompatible with rhs \d+x\d+"): a.cross(b) @@ -1251,7 +1478,7 @@ def test_broadcast_wrong_size_raises(self): """Broadcast other must have the matching flat size (2 or 3).""" a = Matrix(5, 3, [float(i) for i in range(15)]) b = Matrix(1, 2, [1.0, 2.0]) - with pytest.raises(NotImplementedError, + with pytest.raises(ValueError, match=r"cross: .* incompatible with rhs \d+x\d+"): a.cross(b) @@ -1259,7 +1486,7 @@ def test_broadcast_other_must_be_vector_raises(self): """Non-vector other with matching size still rejected (no inferred shape).""" a = Matrix(5, 3, [float(i) for i in range(15)]) b = Matrix(3, 3, [float(i) for i in range(9)]) - with pytest.raises(NotImplementedError, + with pytest.raises(ValueError, match=r"cross: .* incompatible with rhs \d+x\d+"): a.cross(b) @@ -1358,36 +1585,36 @@ def test_axis_none_keyword_matches_default(self): @pytest.mark.parametrize("rows,cols", [(1, 4), (4, 1), (1, 1), (1, 5), (4, 4), (4, 5)]) def test_invalid_shape_raises(self, rows, cols): - """Shapes that aren't 1x2/2x1/Nx2/2xN/1x3/3x1/Nx3/3xN raise NotImplementedError.""" + """Shapes that aren't 1x2/2x1/Nx2/2xN/1x3/3x1/Nx3/3xN raise ValueError.""" n = rows * cols a = Matrix(rows, cols, [float(i) for i in range(n)]) b = Matrix(rows, cols, [float(i + 1) for i in range(n)]) with pytest.raises( - NotImplementedError, + ValueError, match="cross requires a 2D or 3D vector or Nx2 or 2xN or Nx3 or 3xN matrix"): a.cross(b) def test_invalid_axis_raises(self): - """axis outside {None, 0, 1, -1, -2} raises NotImplementedError.""" + """axis outside {None, 0, 1, -1, -2} raises ValueError.""" a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) b = Matrix(2, 2, [5.0, 6.0, 7.0, 8.0]) - with pytest.raises(NotImplementedError, + with pytest.raises(ValueError, match="axis must be -2, -1, 0, or 1"): a.cross(b, axis=2) def test_size_mismatch_scalar_2d_raises(self): - """2D-scalar self with size-3 other raises NotImplementedError.""" + """2D-scalar self with size-3 other raises ValueError.""" a = Matrix(1, 2, [1.0, 2.0]) b = Matrix(1, 3, [3.0, 4.0, 5.0]) - with pytest.raises(NotImplementedError, + with pytest.raises(ValueError, match=r"cross: 2D vector lhs \d+x\d+ incompatible with rhs \d+x\d+"): a.cross(b) def test_size_mismatch_batch_raises(self): - """Batch self with mismatched-shape other raises NotImplementedError.""" + """Batch self with mismatched-shape other raises ValueError.""" a = Matrix(5, 3, [float(i) for i in range(15)]) b = Matrix(7, 3, [float(i) for i in range(21)]) - with pytest.raises(NotImplementedError, + with pytest.raises(ValueError, match=r"cross: Nx3 batch lhs \d+x\d+ incompatible with rhs \d+x\d+"): a.cross(b) @@ -1398,12 +1625,12 @@ def test_in_behavior_propagates_exception(self): @when(a) def result(a, other=other): # noqa: D401 — short behavior - """Trigger a NotImplementedError from inside a behavior.""" + """Trigger a ValueError from inside a behavior.""" a.value.cross(other) wait() assert result.exception is True - assert isinstance(result.value, NotImplementedError) + assert isinstance(result.value, ValueError) assert "cross requires a 2D or 3D vector" in str(result.value) @@ -1491,9 +1718,9 @@ def test_return_value_contract(self, in_place_mode): assert v[0, 2] == 3.0 def test_invalid_axis_raises(self, in_place_mode): - """axis=2 raises NotImplementedError with the canonical message.""" + """axis=2 raises ValueError with the canonical message.""" mat = Matrix(2, 2, [1.0, 0.0, 0.0, 1.0]) - with pytest.raises(NotImplementedError, + with pytest.raises(ValueError, match="axis must be -2, -1, 0, or 1"): mat.normalize(axis=2, in_place=in_place_mode) @@ -1701,14 +1928,14 @@ def test_invalid_shape_raises(self, rows, cols, in_place_mode): """Any shape that is not a 2D vector or N-by-2/2-by-N raises.""" n = rows * cols m = Matrix(rows, cols, [float(i) for i in range(n)]) - with pytest.raises(NotImplementedError, + with pytest.raises(ValueError, match="perpendicular requires a 2D vector or Nx2 or 2xN matrix"): m.perpendicular(in_place=in_place_mode) def test_invalid_axis_raises(self, in_place_mode): - """Out-of-range axis raises NotImplementedError.""" + """Out-of-range axis raises ValueError.""" m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) - with pytest.raises(NotImplementedError, + with pytest.raises(ValueError, match="axis must be -2, -1, 0, or 1"): m.perpendicular(axis=2, in_place=in_place_mode) @@ -1844,17 +2071,17 @@ def test_2x2_axis_minus_two_matches_axis_zero(self): @pytest.mark.parametrize("rows,cols", [(3, 3), (1, 5), (5, 1), (3, 4), (1, 1)]) def test_invalid_shape_raises(self, rows, cols): - """Non 2D-vector shapes raise NotImplementedError.""" + """Non 2D-vector shapes raise ValueError.""" n = rows * cols m = Matrix(rows, cols, [float(i) for i in range(n)]) - with pytest.raises(NotImplementedError, + with pytest.raises(ValueError, match="angle requires a 2D vector or Nx2 or 2xN matrix"): m.angle() def test_invalid_axis_raises(self): - """Out-of-range axis raises NotImplementedError.""" + """Out-of-range axis raises ValueError.""" m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) - with pytest.raises(NotImplementedError, + with pytest.raises(ValueError, match="axis must be -2, -1, 0, or 1"): m.angle(axis=2) @@ -1919,7 +2146,7 @@ def test_overflow_axis_above_long_range(self, factory, method): def test_former_sentinel_no_longer_silent(self, factory, method): """``axis=-1000`` raises rather than being silently treated as a no-axis sentinel.""" mat = factory() - with pytest.raises(NotImplementedError, match="axis must be -2, -1, 0, or 1"): + with pytest.raises(ValueError, match="axis must be -2, -1, 0, or 1"): getattr(mat, method)(axis=-1000) @pytest.mark.parametrize("factory,method", _AXIS_METHODS, @@ -1983,7 +2210,7 @@ def test_cross_axis_accepts_positional(self): @pytest.mark.parametrize("bad_axis,expected_exc,expected_match", [ (True, TypeError, "bool"), (2**32, OverflowError, "out of int range|too large to convert"), - (-1000, NotImplementedError, "axis must be -2, -1, 0, or 1"), + (-1000, ValueError, "axis must be -2, -1, 0, or 1"), ], ids=["bool", "overflow", "sentinel"]) def test_rejected_axis_does_not_mutate_in_place( self, method, bad_axis, expected_exc, expected_match): @@ -2014,7 +2241,7 @@ class TestShapeDisambiguation: def test_1x2_axis0_rejected(self, method): """``1x2`` is row-oriented; ``axis=0`` contradicts and raises.""" m = Matrix(1, 2, [1.0, 2.0]) - with pytest.raises(NotImplementedError, + with pytest.raises(ValueError, match=f"{method} requires a 2D vector or Nx2 or 2xN matrix"): getattr(m, method)(axis=0) @@ -2022,7 +2249,7 @@ def test_1x2_axis0_rejected(self, method): def test_2x1_axis1_rejected(self, method): """``2x1`` is column-oriented; ``axis=1`` contradicts and raises.""" m = Matrix(2, 1, [1.0, 2.0]) - with pytest.raises(NotImplementedError, + with pytest.raises(ValueError, match=f"{method} requires a 2D vector or Nx2 or 2xN matrix"): getattr(m, method)(axis=1) @@ -2030,7 +2257,7 @@ def test_2x1_axis1_rejected(self, method): def test_Nx2_axis0_rejected(self, method): # noqa: N802 (shape name) """``Nx2`` with ``N>2`` is row-oriented; ``axis=0`` contradicts and raises.""" m = Matrix(5, 2, [float(i) for i in range(10)]) - with pytest.raises(NotImplementedError, + with pytest.raises(ValueError, match=f"{method} requires a 2D vector or Nx2 or 2xN matrix"): getattr(m, method)(axis=0) @@ -2038,7 +2265,7 @@ def test_Nx2_axis0_rejected(self, method): # noqa: N802 (shape name) def test_2xN_axis1_rejected(self, method): # noqa: N802 (shape name) """``2xN`` with ``N>2`` is column-oriented; ``axis=1`` contradicts and raises.""" m = Matrix(2, 5, [float(i) for i in range(10)]) - with pytest.raises(NotImplementedError, + with pytest.raises(ValueError, match=f"{method} requires a 2D vector or Nx2 or 2xN matrix"): getattr(m, method)(axis=1) @@ -2046,56 +2273,56 @@ def test_cross_1x2_axis0_rejected(self): """1x2 scalar 2D rejects ``axis=0``.""" a = Matrix(1, 2, [1.0, 2.0]) b = Matrix(1, 2, [3.0, 4.0]) - with pytest.raises(NotImplementedError, match="cross requires a 2D or 3D vector"): + with pytest.raises(ValueError, match="cross requires a 2D or 3D vector"): a.cross(b, axis=0) def test_cross_2x1_axis1_rejected(self): """2x1 scalar 2D rejects ``axis=1``.""" a = Matrix(2, 1, [1.0, 2.0]) b = Matrix(2, 1, [3.0, 4.0]) - with pytest.raises(NotImplementedError, match="cross requires a 2D or 3D vector"): + with pytest.raises(ValueError, match="cross requires a 2D or 3D vector"): a.cross(b, axis=1) def test_cross_1x3_axis0_rejected(self): """1x3 scalar 3D rejects ``axis=0``.""" a = Matrix(1, 3, [1.0, 2.0, 3.0]) b = Matrix(1, 3, [4.0, 5.0, 6.0]) - with pytest.raises(NotImplementedError, match="cross requires a 2D or 3D vector"): + with pytest.raises(ValueError, match="cross requires a 2D or 3D vector"): a.cross(b, axis=0) def test_cross_3x1_axis1_rejected(self): """3x1 scalar 3D rejects ``axis=1``.""" a = Matrix(3, 1, [1.0, 2.0, 3.0]) b = Matrix(3, 1, [4.0, 5.0, 6.0]) - with pytest.raises(NotImplementedError, match="cross requires a 2D or 3D vector"): + with pytest.raises(ValueError, match="cross requires a 2D or 3D vector"): a.cross(b, axis=1) def test_cross_Nx2_axis0_rejected(self): # noqa: N802 (shape name) """Nx2 batch (N>3) is row-oriented; ``axis=0`` contradicts.""" a = Matrix(5, 2, [float(i) for i in range(10)]) b = Matrix(5, 2, [float(i) for i in range(10)]) - with pytest.raises(NotImplementedError, match="cross requires a 2D or 3D vector"): + with pytest.raises(ValueError, match="cross requires a 2D or 3D vector"): a.cross(b, axis=0) def test_cross_2xN_axis1_rejected(self): # noqa: N802 (shape name) """2xN batch (N>3) is column-oriented; ``axis=1`` contradicts.""" a = Matrix(2, 5, [float(i) for i in range(10)]) b = Matrix(2, 5, [float(i) for i in range(10)]) - with pytest.raises(NotImplementedError, match="cross requires a 2D or 3D vector"): + with pytest.raises(ValueError, match="cross requires a 2D or 3D vector"): a.cross(b, axis=1) def test_cross_Nx3_axis0_rejected(self): # noqa: N802 (shape name) """Nx3 batch (N>2) is row-oriented; ``axis=0`` contradicts.""" a = Matrix(5, 3, [float(i) for i in range(15)]) b = Matrix(5, 3, [float(i) for i in range(15)]) - with pytest.raises(NotImplementedError, match="cross requires a 2D or 3D vector"): + with pytest.raises(ValueError, match="cross requires a 2D or 3D vector"): a.cross(b, axis=0) def test_cross_3xN_axis1_rejected(self): # noqa: N802 (shape name) """3xN batch (N>2) is column-oriented; ``axis=1`` contradicts.""" a = Matrix(3, 5, [float(i) for i in range(15)]) b = Matrix(3, 5, [float(i) for i in range(15)]) - with pytest.raises(NotImplementedError, match="cross requires a 2D or 3D vector"): + with pytest.raises(ValueError, match="cross requires a 2D or 3D vector"): a.cross(b, axis=1) def test_2x2_both_axes_accepted(self): @@ -2340,122 +2567,894 @@ def test_transpose_matmul_symmetry(self): assert Matrix.allclose(lhs, rhs) -class TestSelect: - """Tests for Matrix.select() — row and column sub-selection.""" +class TestTake: + """Tests for Matrix.take() — row and column sub-selection.""" - def test_select_rows_with_list(self, mat, shape): - """select([indices], axis=0) returns the requested rows.""" + def test_take_rows_with_list(self, mat, shape): + """take([indices], axis=0) returns the requested rows.""" rows, cols = shape if rows < 2: pytest.skip("need at least 2 rows") indices = [0, rows - 1] - result = mat.select(indices) + result = mat.take(indices) assert result.rows == len(indices) assert result.columns == cols for out_r, src_r in enumerate(indices): for c in range(cols): assert result[out_r, c] == pytest.approx(mat[src_r, c]) - def test_select_rows_with_tuple(self, mat, shape): - """select((indices,), axis=0) also accepts a tuple.""" + def test_take_rows_with_tuple(self, mat, shape): + """take((indices,), axis=0) also accepts a tuple.""" rows, cols = shape if rows < 3: pytest.skip("need at least 3 rows") indices = (1, 0, 2) - result = mat.select(indices) + result = mat.take(indices) assert result.rows == len(indices) for out_r, src_r in enumerate(indices): for c in range(cols): assert result[out_r, c] == pytest.approx(mat[src_r, c]) - def test_select_columns_with_list(self, mat, shape): - """select([indices], axis=1) returns the requested columns.""" + def test_take_columns_with_list(self, mat, shape): + """take([indices], axis=1) returns the requested columns.""" rows, cols = shape if cols < 2: pytest.skip("need at least 2 columns") indices = [cols - 1, 0] - result = mat.select(indices, 1) + result = mat.take(indices, 1) assert result.rows == rows assert result.columns == len(indices) for r in range(rows): for out_c, src_c in enumerate(indices): assert result[r, out_c] == pytest.approx(mat[r, src_c]) - def test_select_columns_with_tuple(self, mat, shape): - """select((indices,), axis=1) also accepts a tuple.""" + def test_take_columns_with_tuple(self, mat, shape): + """take((indices,), axis=1) also accepts a tuple.""" rows, cols = shape if cols < 3: pytest.skip("need at least 3 columns") indices = (2, 0, 1) - result = mat.select(indices, 1) + result = mat.take(indices, 1) assert result.columns == len(indices) for r in range(rows): for out_c, src_c in enumerate(indices): assert result[r, out_c] == pytest.approx(mat[r, src_c]) - def test_select_negative_axis(self, mat, shape): + def test_take_negative_axis(self, mat, shape): """axis=-1 should behave like axis=1 (columns).""" rows, cols = shape if cols < 2: pytest.skip("need at least 2 columns") indices = [0, cols - 1] - result_pos = mat.select(indices, 1) - result_neg = mat.select(indices, -1) + result_pos = mat.take(indices, 1) + result_neg = mat.take(indices, -1) assert Matrix.allclose(result_pos, result_neg) - def test_select_duplicate_indices(self, mat, shape): + def test_take_duplicate_indices(self, mat, shape): """Duplicate indices should duplicate the corresponding rows.""" rows, cols = shape indices = [0, 0, 0] - result = mat.select(indices, 0) + result = mat.take(indices, 0) assert result.rows == 3 for r in range(3): for c in range(cols): assert result[r, c] == pytest.approx(mat[0, c]) - def test_select_single_row(self, mat, shape): - """Selecting a single row returns a 1xcols matrix.""" + def test_take_single_row(self, mat, shape): + """Taking a single row returns a 1xcols matrix.""" rows, cols = shape - result = mat.select([0]) + result = mat.take([0]) assert result.rows == 1 assert result.columns == cols for c in range(cols): assert result[0, c] == pytest.approx(mat[0, c]) - def test_select_single_column(self, mat, shape): - """Selecting a single column returns a rowsx1 matrix.""" + def test_take_single_column(self, mat, shape): + """Taking a single column returns a rowsx1 matrix.""" rows, cols = shape - result = mat.select([0], 1) + result = mat.take([0], 1) assert result.rows == rows assert result.columns == 1 for r in range(rows): assert result[r, 0] == pytest.approx(mat[r, 0]) - def test_select_all_rows_preserves_matrix(self, mat, shape): - """Selecting all rows in order yields an equal matrix.""" + def test_take_all_rows_preserves_matrix(self, mat, shape): + """Taking all rows in order yields an equal matrix.""" rows, cols = shape indices = list(range(rows)) - result = mat.select(indices) + result = mat.take(indices) assert Matrix.allclose(result, mat) - def test_select_all_columns_preserves_matrix(self, mat, shape): - """Selecting all columns in order yields an equal matrix.""" + def test_take_all_columns_preserves_matrix(self, mat, shape): + """Taking all columns in order yields an equal matrix.""" rows, cols = shape indices = list(range(cols)) - result = mat.select(indices, 1) + result = mat.take(indices, 1) assert Matrix.allclose(result, mat) - def test_select_empty_indices_returns_none(self): - """Selecting with an empty list returns None.""" + def test_take_empty_indices_raises(self): + """Taking with an empty list raises IndexError (Option A).""" m = Matrix(3, 3, 1.0) - result = m.select([]) - assert result is None + with pytest.raises(IndexError): + m.take([]) + with pytest.raises(IndexError): + m.take([], 1) + + def test_take_negative_index(self, mat, shape): + """A negative index counts from the end of the axis.""" + rows, cols = shape + last_row = mat.take([-1], 0) + assert last_row.rows == 1 + for c in range(cols): + assert last_row[0, c] == pytest.approx(mat[rows - 1, c]) + last_col = mat.take([-1], 1) + assert last_col.columns == 1 + for r in range(rows): + assert last_col[r, 0] == pytest.approx(mat[r, cols - 1]) + + def test_take_out_of_range_raises(self, mat, shape): + """An out-of-range index raises IndexError naming the value and dim.""" + rows, cols = shape + with pytest.raises(IndexError) as excinfo: + mat.take([rows], 0) + message = str(excinfo.value) + assert str(rows) in message + assert "row" in message + with pytest.raises(IndexError) as excinfo: + mat.take([cols], 1) + message = str(excinfo.value) + assert str(cols) in message + assert "column" in message + + def test_take_negative_boundary(self, mat, shape): + """``-len`` is the first element; ``-len-1`` is out of range.""" + rows, cols = shape + first_row = mat.take([-rows], 0) + for c in range(cols): + assert first_row[0, c] == pytest.approx(mat[0, c]) + with pytest.raises(IndexError): + mat.take([-rows - 1], 0) + first_col = mat.take([-cols], 1) + for r in range(rows): + assert first_col[r, 0] == pytest.approx(mat[r, 0]) + with pytest.raises(IndexError): + mat.take([-cols - 1], 1) - def test_select_invalid_axis_raises(self): + def test_take_overflow_raises(self): + """A value beyond Py_ssize_t raises OverflowError, not IndexError.""" + m = Matrix(3, 3, 1.0) + with pytest.raises(OverflowError): + m.take([10**100], 0) + with pytest.raises(OverflowError): + m.take([10**100], 1) + + def test_take_bool_item(self): + """bool items resolve as 0/1 (bool is an int subclass).""" + m = Matrix(3, 3, 0.0) + for r in range(3): + for c in range(3): + m[r, c] = float(r * 3 + c) + result = m.take([True, False], 0) + assert result.rows == 2 + for c in range(3): + assert result[0, c] == pytest.approx(m[1, c]) + assert result[1, c] == pytest.approx(m[0, c]) + + def test_take_bad_type_raises(self): + """A non-int item raises TypeError.""" + m = Matrix(3, 3, 1.0) + with pytest.raises(TypeError): + m.take(["x"], 0) + + def test_take_int_subscript_parity(self): + """The scalar-int subscript path shares take's boundary outcomes.""" + m = Matrix(3, 3, 1.0) + with pytest.raises(IndexError): + _ = m[-4, 0] + with pytest.raises(OverflowError): + _ = m[10**100, 0] + + def test_take_invalid_axis_raises(self): """axis >= 2 should raise KeyError.""" m = Matrix(3, 3, 1.0) with pytest.raises(KeyError): - m.select([0], 2) + m.take([0], 2) + + def test_take_axis_below_neg2_raises(self): + """An out-of-range negative axis (<= -3) raises KeyError, not a column take.""" + m = Matrix(3, 3, 1.0) + with pytest.raises(KeyError): + m.take([0], -3) + + +class TestFancyIndexing: + """List-key gather through __getitem__ (rows and columns).""" + + def test_getitem_list_rows_matches_take(self, mat, shape): + """m[[r0, r1]] gathers rows, matching take(..., 0).""" + rows, cols = shape + indices = [0, rows - 1] + result = mat[indices] + assert result.rows == len(indices) + assert result.columns == cols + assert Matrix.allclose(result, mat.take(indices, 0)) + + def test_getitem_rows_explicit_full_slice(self, mat, shape): + """m[[r], :] is the same as m[[r]].""" + rows, cols = shape + indices = [0, rows - 1] if rows > 1 else [0] + assert Matrix.allclose(mat[indices, :], mat[indices]) + + def test_getitem_list_cols_matches_take(self, mat, shape): + """m[:, [c0, c1]] gathers columns, matching take(..., 1).""" + rows, cols = shape + indices = [cols - 1, 0] if cols > 1 else [0] + result = mat[:, indices] + assert result.rows == rows + assert result.columns == len(indices) + assert Matrix.allclose(result, mat.take(indices, 1)) + + def test_getitem_negative(self, mat, shape): + """A negative list index counts from the end.""" + rows, cols = shape + result = mat[[-1]] + for c in range(cols): + assert result[0, c] == pytest.approx(mat[rows - 1, c]) + + def test_getitem_duplicates(self, mat, shape): + """Duplicate list indices repeat the row.""" + rows, cols = shape + result = mat[[0, 0, 0]] + assert result.rows == 3 + for r in range(3): + for c in range(cols): + assert result[r, c] == pytest.approx(mat[0, c]) + + def test_getitem_out_of_range(self, mat, shape): + """An out-of-range list index raises IndexError naming value+dim.""" + rows, cols = shape + with pytest.raises(IndexError) as excinfo: + _ = mat[[rows]] + message = str(excinfo.value) + assert str(rows) in message + assert "row" in message + with pytest.raises(IndexError) as excinfo: + _ = mat[:, [cols]] + message = str(excinfo.value) + assert str(cols) in message + assert "column" in message + + def test_getitem_empty_list_raises(self): + """An empty list key raises IndexError (Option A).""" + m = Matrix(3, 3, 1.0) + with pytest.raises(IndexError): + _ = m[[]] + with pytest.raises(IndexError): + _ = m[:, []] + + def test_getitem_tuple_is_2d_not_gather(self): + """m[(0, 0)] returns a float (tuple is 2-D indexing, not gather).""" + m = Matrix(3, 3, 0.0) + m[1, 2] = 7.0 + value = m[(1, 2)] + assert isinstance(value, float) + assert value == pytest.approx(7.0) + + def test_getitem_2d_slice_unchanged(self): + """m[0:2, 1:3] still returns the legacy 2-D sub-matrix.""" + m = Matrix(3, 3, 0.0) + for r in range(3): + for c in range(3): + m[r, c] = float(r * 3 + c) + sub = m[0:2, 1:3] + assert sub.rows == 2 + assert sub.columns == 2 + assert sub[0, 0] == pytest.approx(1.0) + assert sub[1, 1] == pytest.approx(5.0) + + def test_getitem_singleton_list_returns_matrix(self): + """m[[0]] always yields a Matrix; m[0]/m[0,0] yields a float.""" + m = Matrix(1, 1, 3.0) + gathered = m[[0]] + assert isinstance(gathered, Matrix) + assert gathered.rows == 1 + assert gathered.columns == 1 + assert isinstance(m[0, 0], float) + + def test_getitem_paired_list_unsupported(self): + """m[[r], [c]] (paired lists) raises IndexError.""" + m = Matrix(3, 3, 1.0) + with pytest.raises(IndexError): + _ = m[[0, 1], [0, 1]] + + def test_getitem_mixed_list_int_unsupported(self): + """m[[r], c] (list with int) raises IndexError.""" + m = Matrix(3, 3, 1.0) + with pytest.raises(IndexError): + _ = m[[0], 1] + + def test_getitem_partial_slice_list_unsupported(self): + """A list paired with a non-full slice raises IndexError.""" + m = Matrix(3, 3, 1.0) + with pytest.raises(IndexError): + _ = m[1:3, [0]] + with pytest.raises(IndexError): + _ = m[[0], 0:2] + + def test_getitem_bad_type_raises(self): + """A list with a non-int item raises TypeError.""" + m = Matrix(3, 3, 1.0) + with pytest.raises(TypeError): + _ = m[["x"]] + + def test_getitem_boc_roundtrip(self): + """List gather runs inside a @when behavior over a Cown[Matrix].""" + a = Cown(Matrix(3, 2, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0])) + + @when(a) + def result(a): # noqa: D401 — short behavior + """Gather rows inside a behavior and return the result.""" + return a.value[[2, 0]] + + wait() + assert result.exception is False + expected = Matrix(2, 2, [5.0, 6.0, 1.0, 2.0]) + assert Matrix.allclose(result.value, expected) + + def test_getitem_fuzz_permutation(self, mat, shape, rng): + """A random row permutation matches a manually assembled matrix.""" + rows, cols = shape + perm = list(range(rows)) + rng.shuffle(perm) + result = mat[perm] + assert result.rows == rows + assert result.columns == cols + for out_r, src_r in enumerate(perm): + for c in range(cols): + assert result[out_r, c] == pytest.approx(mat[src_r, c]) + + +def _snapshot(m, rows, cols): + """Return a flat list snapshot of a matrix's elements.""" + return [m[r, c] for r in range(rows) for c in range(cols)] + + +class TestScatter: + """List-key scatter assignment through __setitem__ (rows and columns).""" + + def test_scatter_rows_scalar(self): + """m[[rows]] = scalar fills the selected rows.""" + m = Matrix(3, 2, 0.0) + m[[0, 2]] = 5.0 + for c in range(2): + assert m[0, c] == pytest.approx(5.0) + assert m[1, c] == pytest.approx(0.0) + assert m[2, c] == pytest.approx(5.0) + + def test_scatter_rows_matrix(self): + """m[[rows]] = matrix copies each source row (memcpy path).""" + m = Matrix(3, 2, 0.0) + rhs = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + m[[2, 0]] = rhs + assert m[2, 0] == pytest.approx(1.0) + assert m[2, 1] == pytest.approx(2.0) + assert m[0, 0] == pytest.approx(3.0) + assert m[0, 1] == pytest.approx(4.0) + assert m[1, 0] == pytest.approx(0.0) + + def test_scatter_cols_scalar(self): + """m[:, [cols]] = scalar fills the selected columns.""" + m = Matrix(2, 3, 0.0) + m[:, [0, 2]] = 7.0 + for r in range(2): + assert m[r, 0] == pytest.approx(7.0) + assert m[r, 1] == pytest.approx(0.0) + assert m[r, 2] == pytest.approx(7.0) + + def test_scatter_cols_matrix(self): + """m[:, [cols]] = matrix copies each source column (strided).""" + m = Matrix(2, 3, 0.0) + rhs = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + m[:, [2, 0]] = rhs + assert m[0, 2] == pytest.approx(1.0) + assert m[0, 0] == pytest.approx(2.0) + assert m[1, 2] == pytest.approx(3.0) + assert m[1, 0] == pytest.approx(4.0) + assert m[0, 1] == pytest.approx(0.0) + + def test_scatter_rows_explicit_slice(self): + """m[[rows], :] = v is the same as m[[rows]] = v.""" + a = Matrix(3, 2, 0.0) + b = Matrix(3, 2, 0.0) + a[[0, 1]] = 4.0 + b[[0, 1], :] = 4.0 + assert Matrix.allclose(a, b) + + def test_scatter_then_gather_roundtrip(self, mat, shape): + """A scatter then a gather of the same rows returns the RHS.""" + rows, cols = shape + indices = [0, rows - 1] if rows > 1 else [0] + rhs = Matrix(len(indices), cols, + [float(i) for i in range(len(indices) * cols)]) + mat[indices] = rhs + assert Matrix.allclose(mat[indices], rhs) + + def test_scatter_negative_index(self): + """A negative list index assigns from the end.""" + m = Matrix(3, 2, 0.0) + m[[-1]] = 9.0 + for c in range(2): + assert m[2, c] == pytest.approx(9.0) + assert m[0, c] == pytest.approx(0.0) + + def test_scatter_oob_message_matches_gather(self): + """Read and write OOB share the resolver and the message.""" + m = Matrix(3, 3, 1.0) + with pytest.raises(IndexError) as read_exc: + _ = m[[3]] + with pytest.raises(IndexError) as write_exc: + m[[3]] = 1.0 + assert str(read_exc.value) == str(write_exc.value) + + def test_scatter_duplicate_last_write_wins(self): + """m[[1, 1]] = matrix leaves row 1 equal to the last source row.""" + m = Matrix(2, 2, 0.0) + rhs = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + m[[1, 1]] = rhs + assert m[1, 0] == pytest.approx(3.0) + assert m[1, 1] == pytest.approx(4.0) + + def test_scatter_iadd_duplicate_adds_once(self): + """m[[0, 0]] += v increments row 0 once (last-write-wins).""" + m = Matrix(2, 2, [1.0, 1.0, 1.0, 1.0]) + m[[0, 0]] += 5.0 + for c in range(2): + assert m[0, c] == pytest.approx(6.0) + assert m[1, c] == pytest.approx(1.0) + + def test_iadd_whole_matrix_sanity(self): + """m += scalar works (the slot the augmented scatter relies on).""" + m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + m += 1.0 + assert Matrix.allclose(m, Matrix(2, 2, [2.0, 3.0, 4.0, 5.0])) + + def test_scatter_augmented_ops_rows(self): + """+= -= *= /= on a row selection match gather + manual op.""" + for op, fn in ( + ("+=", lambda a, b: a + b), + ("-=", lambda a, b: a - b), + ("*=", lambda a, b: a * b), + ("/=", lambda a, b: a / b), + ): + m = Matrix(3, 2, [float(v) for v in range(1, 7)]) + expected = [ + [fn(m[r, c], 2.0) for c in range(2)] if r in (0, 2) + else [m[r, c] for c in range(2)] + for r in range(3) + ] + if op == "+=": + m[[0, 2]] += 2.0 + elif op == "-=": + m[[0, 2]] -= 2.0 + elif op == "*=": + m[[0, 2]] *= 2.0 + else: + m[[0, 2]] /= 2.0 + for r in range(3): + for c in range(2): + assert m[r, c] == pytest.approx(expected[r][c]) + + def test_scatter_augmented_op_cols(self): + """m[:, [c]] *= v matches gather + manual op on that column.""" + m = Matrix(2, 3, [float(v) for v in range(1, 7)]) + expected = [[m[r, c] * 3.0 if c == 1 else m[r, c] + for c in range(3)] for r in range(2)] + m[:, [1]] *= 3.0 + for r in range(2): + for c in range(3): + assert m[r, c] == pytest.approx(expected[r][c]) + + def test_scatter_oob_leaves_matrix_unmodified(self): + """A bad index found in phase 1 aborts before any write.""" + m = Matrix(3, 3, 1.0) + before = _snapshot(m, 3, 3) + with pytest.raises(IndexError): + m[[0, 99]] = 5.0 + assert _snapshot(m, 3, 3) == before + + def test_scatter_shape_mismatch_leaves_unmodified(self): + """A wrong-shaped RHS found in phase 2 aborts before any write.""" + m = Matrix(3, 3, 1.0) + before = _snapshot(m, 3, 3) + with pytest.raises(ValueError): + m[[0, 1]] = Matrix(3, 3, 2.0) + assert _snapshot(m, 3, 3) == before + + def test_scatter_shape_mismatch_message(self): + """The shape-mismatch ValueError names both shapes.""" + m = Matrix(4, 3, 1.0) + with pytest.raises(ValueError) as excinfo: + m[[0, 1]] = Matrix(3, 3, 2.0) + message = str(excinfo.value) + assert "3x3" in message + assert "2x3" in message + + @pytest.mark.parametrize("count", [64, 65, 100]) + def test_scatter_count_64_65_100_rows(self, count): + """Row scatter exercises the stack/heap index-buffer boundary.""" + m = Matrix(3, 2, 0.0) + idx = [1] * count + rhs = Matrix(count, 2, [float(i) for i in range(count * 2)]) + m[idx] = rhs + assert m[1, 0] == pytest.approx(float((count - 1) * 2)) + assert m[1, 1] == pytest.approx(float((count - 1) * 2 + 1)) + for c in range(2): + assert m[0, c] == pytest.approx(0.0) + assert m[2, c] == pytest.approx(0.0) + + @pytest.mark.parametrize("count", [64, 65, 100]) + def test_scatter_count_64_65_100_cols(self, count): + """Column scatter exercises the stack/heap index-buffer boundary.""" + m = Matrix(2, 3, 0.0) + idx = [1] * count + rhs = Matrix(2, count, [float(i) for i in range(2 * count)]) + m[:, idx] = rhs + assert m[0, 1] == pytest.approx(float(count - 1)) + assert m[1, 1] == pytest.approx(float(2 * count - 1)) + for r in range(2): + assert m[r, 0] == pytest.approx(0.0) + assert m[r, 2] == pytest.approx(0.0) + + def test_scatter_non_int_item(self): + """A non-int list item raises TypeError (incl. heap-sized lists).""" + m = Matrix(3, 3, 1.0) + with pytest.raises(TypeError): + m[[0, 1.5]] = 1.0 + with pytest.raises(TypeError): + m[[0, None]] = 1.0 + with pytest.raises(TypeError): + m[[0] * 70 + [1.5]] = 1.0 + + def test_scatter_bool_item(self): + """Bool list items are accepted as 1/0.""" + m = Matrix(2, 2, 0.0) + m[[True, False]] = 3.0 + for c in range(2): + assert m[0, c] == pytest.approx(3.0) + assert m[1, c] == pytest.approx(3.0) + + def test_scatter_cols_single_row(self): + """Column scatter degenerates correctly when R == 1.""" + m = Matrix(1, 3, 0.0) + m[:, [0, 2]] = Matrix(1, 2, [4.0, 5.0]) + assert m[0, 0] == pytest.approx(4.0) + assert m[0, 2] == pytest.approx(5.0) + assert m[0, 1] == pytest.approx(0.0) + + def test_scatter_1x1_matrix_as_scalar(self): + """A 1x1 matrix RHS broadcasts like a scalar.""" + m = Matrix(3, 2, 0.0) + m[[0, 1]] = Matrix(1, 1, 8.0) + for c in range(2): + assert m[0, c] == pytest.approx(8.0) + assert m[1, c] == pytest.approx(8.0) + + def test_scatter_empty_list_raises(self): + """An empty list key raises IndexError (Option A).""" + m = Matrix(3, 3, 1.0) + with pytest.raises(IndexError): + m[[]] = 1.0 + with pytest.raises(IndexError): + m[:, []] = 1.0 + + def test_scatter_paired_list_unsupported(self): + """m[[r], [c]] = v raises IndexError.""" + m = Matrix(3, 3, 1.0) + with pytest.raises(IndexError): + m[[0, 1], [0, 1]] = 1.0 + + def test_scatter_mixed_list_int_unsupported(self): + """m[[r], c] = v raises IndexError.""" + m = Matrix(3, 3, 1.0) + with pytest.raises(IndexError): + m[[0], 1] = 1.0 + + def test_scatter_partial_slice_list_unsupported(self): + """A list paired with a non-full slice raises IndexError.""" + m = Matrix(3, 3, 1.0) + with pytest.raises(IndexError): + m[1:3, [0]] = 1.0 + with pytest.raises(IndexError): + m[[0], 0:2] = 1.0 + + def test_scatter_bad_rhs_type(self): + """A non-numeric, non-matrix RHS raises TypeError.""" + m = Matrix(3, 3, 1.0) + with pytest.raises(TypeError): + m[[0]] = "x" + + def test_scatter_2d_assignment_unchanged(self): + """Legacy int/slice/tuple assignment is unaffected.""" + m = Matrix(3, 3, 0.0) + m[1, 2] = 5.0 + assert m[1, 2] == pytest.approx(5.0) + m[0:2, 0:2] = 7.0 + for r in range(2): + for c in range(2): + assert m[r, c] == pytest.approx(7.0) + + def test_scatter_not_acquired_raises(self): + """Scatter on a cown-resident (unacquired) matrix raises.""" + m = Matrix(3, 2, 1.0) + Cown(m) + with pytest.raises(RuntimeError): + m[[0]] = 1.0 + + def test_scatter_rhs_not_acquired_raises(self): + """A cown-resident RHS matrix raises (ownership gate).""" + m = Matrix(3, 2, 1.0) + rhs = Matrix(1, 2, 2.0) + Cown(rhs) + with pytest.raises(RuntimeError): + m[[0]] = rhs + + def test_scatter_boc_roundtrip(self): + """Scatter runs inside a @when behavior over a Cown[Matrix].""" + a = Cown(Matrix(3, 2, 0.0)) + + @when(a) + def result(a): # noqa: D401 — short behavior + """Scatter rows and a column inside a behavior.""" + a.value[[0, 2]] = 5.0 + a.value[:, [1]] = Matrix(3, 1, [7.0, 8.0, 9.0]) + return a.value[[0, 1, 2]] + + wait() + assert result.exception is False + expected = Matrix(3, 2, [5.0, 7.0, 0.0, 8.0, 5.0, 9.0]) + assert Matrix.allclose(result.value, expected) + + def test_scatter_fuzz(self, mat, shape, rng): + """A random permutation scatter then gather returns the RHS.""" + rows, cols = shape + perm = list(range(rows)) + rng.shuffle(perm) + rhs = Matrix(rows, cols, + [rng.uniform(-50, 50) for _ in range(rows * cols)]) + mat[perm] = rhs + assert Matrix.allclose(mat[perm], rhs) + + def test_scatter_self_alias_swaps_rows(self): + """A self-aliased row scatter permutes from pre-write values.""" + m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + m[[1, 0]] = m + assert Matrix.allclose(m, Matrix(2, 2, [3.0, 4.0, 1.0, 2.0])) + + def test_scatter_self_alias_swaps_cols(self): + """A self-aliased column scatter permutes from pre-write values.""" + m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + m[:, [1, 0]] = m + assert Matrix.allclose(m, Matrix(2, 2, [2.0, 1.0, 4.0, 3.0])) + + def test_scatter_self_alias_identity_unchanged(self): + """A self-aliased identity scatter leaves the matrix unchanged.""" + m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + m[[0, 1]] = m + assert Matrix.allclose(m, Matrix(2, 2, [1.0, 2.0, 3.0, 4.0])) + + +class TestPut: + """Explicit put() method — the write-side counterpart of take().""" + + def test_put_rows_scalar(self): + """put(rows, scalar) fills the selected rows.""" + m = Matrix(3, 2, 0.0) + m.put([0, 2], 5.0) + for c in range(2): + assert m[0, c] == pytest.approx(5.0) + assert m[1, c] == pytest.approx(0.0) + assert m[2, c] == pytest.approx(5.0) + + def test_put_rows_matrix(self): + """put(rows, matrix) copies each source row.""" + m = Matrix(3, 2, 0.0) + rhs = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + m.put([2, 0], rhs) + assert m[2, 0] == pytest.approx(1.0) + assert m[0, 0] == pytest.approx(3.0) + + def test_put_cols_matrix(self): + """put(cols, matrix, axis=1) copies each source column.""" + m = Matrix(2, 3, 0.0) + rhs = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + m.put([2, 0], rhs, axis=1) + assert m[0, 2] == pytest.approx(1.0) + assert m[0, 0] == pytest.approx(2.0) + assert m[1, 2] == pytest.approx(3.0) + assert m[1, 0] == pytest.approx(4.0) + + def test_put_matches_subscript_rows(self, mat, shape): + """put(rows, v) equals m[rows] = v.""" + rows, cols = shape + indices = [0, rows - 1] if rows > 1 else [0] + a = mat.copy() + b = mat.copy() + rhs = Matrix(len(indices), cols, + [float(i) for i in range(len(indices) * cols)]) + a.put(indices, rhs) + b[indices] = rhs + assert Matrix.allclose(a, b) + + def test_put_matches_subscript_cols(self, mat, shape): + """put(cols, v, axis=1) equals m[:, cols] = v.""" + rows, cols = shape + indices = [0, cols - 1] if cols > 1 else [0] + a = mat.copy() + b = mat.copy() + rhs = Matrix(rows, len(indices), + [float(i) for i in range(rows * len(indices))]) + a.put(indices, rhs, axis=1) + b[:, indices] = rhs + assert Matrix.allclose(a, b) + + def test_put_returns_self(self): + """put returns self to allow chaining.""" + m = Matrix(3, 2, 0.0) + result = m.put([0], 1.0) + assert result is m + + def test_put_negative_axis(self): + """A negative axis maps -1 to columns.""" + m = Matrix(2, 3, 0.0) + m.put([0], 9.0, axis=-1) + for r in range(2): + assert m[r, 0] == pytest.approx(9.0) + assert m[r, 1] == pytest.approx(0.0) + + def test_put_negative_index(self): + """A negative row index assigns from the end.""" + m = Matrix(3, 2, 0.0) + m.put([-1], 7.0) + for c in range(2): + assert m[2, c] == pytest.approx(7.0) + + def test_put_1x1_as_scalar(self): + """A 1x1 matrix value broadcasts like a scalar.""" + m = Matrix(3, 2, 0.0) + m.put([0, 1], Matrix(1, 1, 8.0)) + for c in range(2): + assert m[0, c] == pytest.approx(8.0) + assert m[1, c] == pytest.approx(8.0) + + def test_put_empty_raises(self): + """An empty index list raises IndexError.""" + m = Matrix(3, 3, 1.0) + with pytest.raises(IndexError): + m.put([], 1.0) + + def test_put_out_of_range_raises(self): + """An out-of-range index raises IndexError and writes nothing.""" + m = Matrix(3, 3, 1.0) + before = _snapshot(m, 3, 3) + with pytest.raises(IndexError): + m.put([0, 99], 5.0) + assert _snapshot(m, 3, 3) == before + + def test_put_shape_mismatch_raises(self): + """A wrong-shaped matrix value raises ValueError and writes nothing.""" + m = Matrix(3, 3, 1.0) + before = _snapshot(m, 3, 3) + with pytest.raises(ValueError): + m.put([0, 1], Matrix(3, 3, 2.0)) + assert _snapshot(m, 3, 3) == before + + def test_put_invalid_axis_raises(self): + """axis >= 2 raises KeyError, matching take().""" + m = Matrix(3, 3, 1.0) + with pytest.raises(KeyError): + m.put([0], 1.0, 2) + + def test_put_axis_below_neg2_raises(self): + """An out-of-range negative axis (<= -3) raises KeyError, not a column put.""" + m = Matrix(3, 3, 1.0) + with pytest.raises(KeyError): + m.put([0], 1.0, -3) + + def test_put_bad_value_type_raises(self): + """A non-numeric, non-matrix value raises TypeError.""" + m = Matrix(3, 3, 1.0) + with pytest.raises(TypeError): + m.put([0], "x") + + def test_put_accumulate_scalar_rows(self): + """accumulate=True folds a scalar additively over duplicate rows.""" + m = Matrix(3, 2, 1.0) + m.put([0, 0, 2], 2.0, accumulate=True) + assert Matrix.allclose( + m, Matrix(3, 2, [5.0, 5.0, 1.0, 1.0, 3.0, 3.0])) + + def test_put_accumulate_matrix_rows(self): + """accumulate=True adds matching value rows into duplicate rows.""" + m = Matrix(2, 2, 0.0) + v = Matrix(3, 2, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + m.put([0, 0, 1], v, accumulate=True) + # row 0 gets v rows 0 and 1; row 1 gets v row 2. + assert Matrix.allclose(m, Matrix(2, 2, [4.0, 6.0, 5.0, 6.0])) + + def test_put_accumulate_scalar_cols(self): + """accumulate=True folds a scalar additively over duplicate columns.""" + m = Matrix(2, 3, 1.0) + m.put([1, 1], 2.0, axis=1, accumulate=True) + assert Matrix.allclose( + m, Matrix(2, 3, [1.0, 5.0, 1.0, 1.0, 5.0, 1.0])) + + def test_put_accumulate_matrix_cols(self): + """accumulate=True adds matching value columns into duplicate cols.""" + m = Matrix(2, 2, 0.0) + v = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + m.put([0, 0, 1], v, axis=1, accumulate=True) + # col 0 gets v cols 0 and 1; col 1 gets v col 2. + assert Matrix.allclose(m, Matrix(2, 2, [3.0, 3.0, 9.0, 6.0])) + + def test_put_accumulate_default_is_last_write_wins(self): + """Without accumulate, duplicate indices remain last-write-wins.""" + m = Matrix(3, 2, 1.0) + m.put([0, 0], 2.0) + assert Matrix.allclose( + m, Matrix(3, 2, [2.0, 2.0, 1.0, 1.0, 1.0, 1.0])) + + def test_put_accumulate_self_alias(self): + """A self-aliased accumulate reads pre-write values via the snapshot.""" + m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + m.put([0, 1], m, accumulate=True) + # Each row doubles; the snapshot keeps reads stable under aliasing. + assert Matrix.allclose(m, Matrix(2, 2, [2.0, 4.0, 6.0, 8.0])) + + def test_put_accumulate_self_alias_cols(self): + """A self-aliased column accumulate reads stable pre-write columns.""" + m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + m.put([0, 1], m, axis=1, accumulate=True) + # Each column doubles; the snapshot keeps reads stable under aliasing. + assert Matrix.allclose(m, Matrix(2, 2, [2.0, 4.0, 6.0, 8.0])) + + def test_put_accumulate_self_alias_1x1(self): + """A 1x1 self-slice accumulates as a scalar (captured pre-write).""" + m = Matrix(1, 1, [3.0]) + m.put([0], m, accumulate=True) + # The 1x1 RHS is read as a scalar before the write, so this doubles. + assert Matrix.allclose(m, Matrix(1, 1, [6.0])) + + def test_put_accumulate_order_preserving(self): + """Folds happen in index-list order, giving a bit-exact result.""" + m = Matrix(1, 1, [0.0]) + # Left-to-right: ((0 + 1e16) + 1.0) + -1e16 == 0.0, but a reordered + # or contraction-fused fold would not cancel exactly. + v = Matrix(3, 1, [1e16, 1.0, -1e16]) + m.put([0, 0, 0], v, accumulate=True) + assert m[0, 0] == 0.0 + + def test_put_not_acquired_raises(self): + """put on a cown-resident matrix raises.""" + m = Matrix(3, 2, 1.0) + Cown(m) + with pytest.raises(RuntimeError): + m.put([0], 1.0) + + def test_put_boc_roundtrip(self): + """put runs inside a @when behavior over a Cown[Matrix].""" + a = Cown(Matrix(3, 2, 0.0)) + + @when(a) + def result(a): # noqa: D401 — short behavior + """Put rows and a column inside a behavior.""" + a.value.put([0, 2], 5.0) + a.value.put([1], Matrix(3, 1, [7.0, 8.0, 9.0]), axis=1) + return a.value.copy() + + wait() + assert result.exception is False + expected = Matrix(3, 2, [5.0, 7.0, 0.0, 8.0, 5.0, 9.0]) + assert Matrix.allclose(result.value, expected) VECTOR_LENGTHS = [1, 3, 5, 10, 32] @@ -2669,7 +3668,7 @@ def test_broadcast_mismatch_raises(self): """Mismatched non-broadcastable shapes raise an error.""" a = Matrix(3, 4) b = Matrix(2, 5) - with pytest.raises(NotImplementedError): + with pytest.raises(ValueError): _ = a + b @@ -2980,7 +3979,7 @@ class TestClip: """Tests for the clip() method.""" def test_clip_two_args(self, mat, shape, random_values): - """clip(minval, maxval) clamps every element.""" + """clip(min, max) clamps every element to [min, max].""" rows, cols = shape lo, hi = -25.0, 25.0 c = mat.clip(lo, hi) @@ -2992,18 +3991,43 @@ def test_clip_two_args(self, mat, shape, random_values): expected = max(lo, min(hi, v)) assert c[i, j] == pytest.approx(expected) - def test_clip_one_arg(self, shape, rng): - """clip(maxval) clamps to [0, maxval].""" + def test_clip_min_only(self, shape, rng): + """clip(min=lo) clamps only below, leaving the upper end unbounded.""" + rows, cols = shape + vals = [rng.uniform(-10, 10) for _ in range(rows * cols)] + m = Matrix(rows, cols, vals) + lo = -2.0 + c = m.clip(min=lo) + for i in range(rows): + for j in range(cols): + v = vals[i * cols + j] + assert c[i, j] == pytest.approx(max(lo, v)) + + def test_clip_max_only(self, shape, rng): + """clip(max=hi) clamps only above, leaving the lower end unbounded.""" rows, cols = shape vals = [rng.uniform(-10, 10) for _ in range(rows * cols)] m = Matrix(rows, cols, vals) hi = 5.0 - c = m.clip(hi) + c = m.clip(max=hi) for i in range(rows): for j in range(cols): v = vals[i * cols + j] - expected = max(0.0, min(hi, v)) - assert c[i, j] == pytest.approx(expected) + assert c[i, j] == pytest.approx(min(hi, v)) + + def test_clip_first_positional_is_min(self, shape): + """The single positional argument is the lower bound (NumPy order).""" + rows, cols = shape + m = Matrix(rows, cols, -5.0) + c = m.clip(0.0) + expected = Matrix(rows, cols, 0.0) + assert Matrix.allclose(c, expected) + + def test_clip_no_bounds_raises(self): + """clip() with neither bound raises ValueError.""" + m = Matrix(2, 2, 1.0) + with pytest.raises(ValueError): + m.clip() def test_clip_no_change(self, shape): """Values already within range are unchanged.""" @@ -3030,7 +4054,7 @@ def test_clip_all_above(self, shape): assert Matrix.allclose(c, expected) def test_clip_invalid_range(self): - """clip() raises AssertionError when maxval < minval.""" + """clip() raises AssertionError when max < min.""" m = Matrix(2, 2, 1.0) with pytest.raises(AssertionError): m.clip(10.0, 0.0) @@ -3237,6 +4261,11 @@ def test_vector_as_column_false_is_row(self): for i in range(3): assert v[0, i] == pytest.approx(vals[i]) + def test_vector_empty_raises(self): + """Matrix.vector([]) cannot form a matrix and raises ValueError.""" + with pytest.raises(ValueError): + Matrix.vector([]) + class TestConcat: """Tests for Matrix.concat() — concatenation along rows or columns.""" @@ -3492,9 +4521,9 @@ def test_argextreme_axis1_fuzz(self, mat, shape, random_values, want_max): assert result[r, 0] == _flat_argextreme(row, want_max) def test_argmin_invalid_axis_raises(self): - """An out-of-range axis raises NotImplementedError.""" + """An out-of-range axis raises ValueError.""" m = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) - with pytest.raises(NotImplementedError): + with pytest.raises(ValueError): m.argmin(axis=2) def test_argextreme_nan_in_middle_is_skipped(self): @@ -3597,14 +4626,14 @@ def test_outer_inplace_raises(self): """In-place outer broadcast would change shape and is rejected.""" col = Matrix(3, 1, [1.0, 2.0, 3.0]) row = Matrix(1, 4, [1.0, 2.0, 3.0, 4.0]) - with pytest.raises(NotImplementedError): + with pytest.raises(ValueError): col *= row def test_incompatible_shapes_still_raise(self): """A row vector whose width mismatches a full matrix still raises.""" a = Matrix(1, 3, [1.0, 2.0, 3.0]) b = Matrix(2, 4, 1.0) - with pytest.raises(NotImplementedError): + with pytest.raises(ValueError): _ = a * b @@ -4403,7 +5432,7 @@ def test_matmul_incompatible_shapes_raises(self): """@ raises when inner dimensions don't match.""" a = Matrix(2, 3) b = Matrix(4, 2) - with pytest.raises(NotImplementedError): + with pytest.raises(ValueError): _ = a @ b @@ -4783,3 +5812,709 @@ def check(v): assert r0 == pytest.approx(-1.0) assert r1 == pytest.approx(2.0) assert r2 == pytest.approx(-3.0) + + +def _flat(m): + """Row-major flat list of a matrix's elements.""" + rows, cols = m.rows, m.columns + return [m[i, j] for i in range(rows) for j in range(cols)] + + +class TestSqrt: + """Element-wise square root.""" + + def test_sqrt_matches_math(self, shape, rng): + """sqrt() matches math.sqrt element-wise for non-negative inputs.""" + rows, cols = shape + vals = [rng.uniform(0, 100) for _ in range(rows * cols)] + m = Matrix(rows, cols, vals) + result = m.sqrt() + for i in range(rows): + for j in range(cols): + assert result[i, j] == pytest.approx(math.sqrt(vals[i * cols + j])) + + def test_sqrt_negative_is_nan(self): + """Negative elements yield NaN rather than raising.""" + m = Matrix(1, 3, [-1.0, -4.0, 9.0]) + result = m.sqrt() + assert math.isnan(result[0, 0]) + assert math.isnan(result[0, 1]) + assert result[0, 2] == pytest.approx(3.0) + + def test_sqrt_zero(self): + """sqrt(0) is 0.""" + m = Matrix(1, 1, [0.0]) + assert m.sqrt()[0, 0] == pytest.approx(0.0) + + def test_sqrt_out_of_place_preserves_source(self, shape, rng): + """Default sqrt() returns a new matrix and leaves the source intact.""" + rows, cols = shape + vals = [rng.uniform(0, 100) for _ in range(rows * cols)] + m = Matrix(rows, cols, vals) + _ = m.sqrt() + for i in range(rows): + for j in range(cols): + assert m[i, j] == pytest.approx(vals[i * cols + j]) + + def test_sqrt_in_place(self, shape, rng): + """sqrt(in_place=True) mutates self and returns it.""" + rows, cols = shape + vals = [rng.uniform(0, 100) for _ in range(rows * cols)] + m = Matrix(rows, cols, vals) + result = m.sqrt(in_place=True) + assert result is m + for i in range(rows): + for j in range(cols): + assert m[i, j] == pytest.approx(math.sqrt(vals[i * cols + j])) + + +COMPARE_OPS = [ + ("less", lambda a, b: a < b), + ("less_equal", lambda a, b: a <= b), + ("greater", lambda a, b: a > b), + ("greater_equal", lambda a, b: a >= b), + ("equal", lambda a, b: a == b), + ("not_equal", lambda a, b: a != b), +] + + +class TestComparisonMasks: + """Element-wise comparison methods producing 0/1 mask matrices.""" + + @pytest.mark.parametrize("name,ref", COMPARE_OPS, ids=[o[0] for o in COMPARE_OPS]) + def test_same_shape(self, name, ref, shape, rng): + """Each mask method matches Python's element-wise comparison.""" + rows, cols = shape + # Overlapping integer-valued draws so ties exercise == / <= / >=. + a_vals = [float(rng.randint(0, 5)) for _ in range(rows * cols)] + b_vals = [float(rng.randint(0, 5)) for _ in range(rows * cols)] + a = Matrix(rows, cols, a_vals) + b = Matrix(rows, cols, b_vals) + result = getattr(a, name)(b) + assert result.rows == rows + assert result.columns == cols + for i in range(rows): + for j in range(cols): + expected = 1.0 if ref(a_vals[i * cols + j], b_vals[i * cols + j]) else 0.0 + assert result[i, j] == expected + + @pytest.mark.parametrize("name,ref", COMPARE_OPS, ids=[o[0] for o in COMPARE_OPS]) + def test_scalar(self, name, ref, shape, rng): + """Mask methods broadcast a scalar against every element.""" + rows, cols = shape + a_vals = [float(rng.randint(0, 5)) for _ in range(rows * cols)] + a = Matrix(rows, cols, a_vals) + scalar = 3.0 + result = getattr(a, name)(scalar) + for i in range(rows): + for j in range(cols): + expected = 1.0 if ref(a_vals[i * cols + j], scalar) else 0.0 + assert result[i, j] == expected + + @pytest.mark.parametrize("name,ref", COMPARE_OPS, ids=[o[0] for o in COMPARE_OPS]) + def test_row_vector_broadcast(self, name, ref): + """Mask methods broadcast a 1xN row vector across rows.""" + a = Matrix(2, 3, [1.0, 5.0, 3.0, 4.0, 2.0, 3.0]) + v = Matrix(1, 3, [3.0, 3.0, 3.0]) + result = getattr(a, name)(v) + for i in range(2): + for j in range(3): + expected = 1.0 if ref(a[i, j], v[0, j]) else 0.0 + assert result[i, j] == expected + + @pytest.mark.parametrize("name,ref", COMPARE_OPS, ids=[o[0] for o in COMPARE_OPS]) + def test_column_vector_broadcast(self, name, ref): + """Mask methods broadcast an Mx1 column vector across columns.""" + a = Matrix(3, 2, [1.0, 5.0, 3.0, 4.0, 2.0, 3.0]) + v = Matrix(3, 1, [3.0, 3.0, 3.0]) + result = getattr(a, name)(v) + for i in range(3): + for j in range(2): + expected = 1.0 if ref(a[i, j], v[i, 0]) else 0.0 + assert result[i, j] == expected + + @pytest.mark.parametrize("name,ref", COMPARE_OPS, ids=[o[0] for o in COMPARE_OPS]) + def test_one_by_one_broadcast(self, name, ref, shape, rng): + """A 1x1 matrix broadcasts like a scalar in mask methods.""" + rows, cols = shape + a_vals = [float(rng.randint(0, 5)) for _ in range(rows * cols)] + a = Matrix(rows, cols, a_vals) + scalar = 3.0 + result = getattr(a, name)(Matrix(1, 1, [scalar])) + for i in range(rows): + for j in range(cols): + expected = 1.0 if ref(a_vals[i * cols + j], scalar) else 0.0 + assert result[i, j] == expected + + def test_nan_less_is_zero(self): + """NaN comparisons follow IEEE: ordered ops yield 0, not_equal yields 1.""" + nan = float("nan") + a = Matrix(1, 2, [nan, 1.0]) + assert a.less(0.0)[0, 0] == 0.0 + assert a.less_equal(0.0)[0, 0] == 0.0 + assert a.greater(0.0)[0, 0] == 0.0 + assert a.greater_equal(0.0)[0, 0] == 0.0 + assert a.equal(nan)[0, 0] == 0.0 + assert a.not_equal(nan)[0, 0] == 1.0 + + def test_masks_are_zero_or_one(self, shape, rng): + """Every mask element is exactly 0.0 or 1.0.""" + rows, cols = shape + a = Matrix(rows, cols, [rng.uniform(-9, 9) for _ in range(rows * cols)]) + b = Matrix(rows, cols, [rng.uniform(-9, 9) for _ in range(rows * cols)]) + for name, _ in COMPARE_OPS: + for v in _flat(getattr(a, name)(b)): + assert v in (0.0, 1.0) + + +RICHCOMPARE_OPS = [ + ("lt", lambda a, b: a < b), + ("le", lambda a, b: a <= b), + ("gt", lambda a, b: a > b), + ("ge", lambda a, b: a >= b), + ("eq", lambda a, b: a == b), + ("ne", lambda a, b: a != b), +] + + +class TestLexicographicCompare: + """Operators (<, <=, >, >=, ==, !=) return a single lexicographic bool.""" + + @pytest.mark.parametrize("name,op", RICHCOMPARE_OPS, ids=[o[0] for o in RICHCOMPARE_OPS]) + def test_matches_list_ordering(self, name, op, shape, rng): + """Operator result matches Python list comparison of flat elements.""" + rows, cols = shape + a_vals = [float(rng.randint(0, 3)) for _ in range(rows * cols)] + b_vals = [float(rng.randint(0, 3)) for _ in range(rows * cols)] + a = Matrix(rows, cols, a_vals) + b = Matrix(rows, cols, b_vals) + assert op(a, b) is op(a_vals, b_vals) + + def test_equal_same_values(self, shape, rng): + """Equal-valued same-shape matrices compare equal.""" + rows, cols = shape + vals = [rng.uniform(-9, 9) for _ in range(rows * cols)] + a = Matrix(rows, cols, list(vals)) + b = Matrix(rows, cols, list(vals)) + assert a == b + assert not (a != b) + assert a <= b + assert a >= b + assert not (a < b) + assert not (a > b) + + def test_first_difference_decides(self): + """The first differing element decides, later elements ignored.""" + a = Matrix(1, 3, [1.0, 9.0, 9.0]) + b = Matrix(1, 3, [2.0, 0.0, 0.0]) + assert a < b + assert not (a > b) + + @pytest.mark.parametrize("name,op", RICHCOMPARE_OPS, ids=[o[0] for o in RICHCOMPARE_OPS]) + def test_scalar_lexicographic(self, name, op): + """Scalar comparison broadcasts and decides on element 0 unless tied.""" + a = Matrix(1, 3, [5.0, 1.0, 1.0]) + assert op(a, 5.0) is op([5.0, 1.0, 1.0], [5.0, 5.0, 5.0]) + + def test_eq_ne_shape_mismatch_no_raise(self): + """== / != return False / True on a shape mismatch without raising.""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + c = Matrix(1, 4, [1.0, 2.0, 3.0, 4.0]) + assert (a == c) is False + assert (a != c) is True + + def test_eq_same_values_different_shape_not_equal(self): + """Same flat data but different shape is not equal.""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + c = Matrix(1, 4, [1.0, 2.0, 3.0, 4.0]) + assert a != c + + @pytest.mark.parametrize("op", [ + lambda a, b: a < b, + lambda a, b: a <= b, + lambda a, b: a > b, + lambda a, b: a >= b, + ]) + def test_ordering_shape_mismatch_raises(self, op): + """Ordering operators raise ValueError on a matrix shape mismatch.""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + c = Matrix(1, 4, [1.0, 2.0, 3.0, 4.0]) + with pytest.raises(ValueError): + op(a, c) + + def test_eq_non_matrix_is_false(self): + """== against a non-matrix, non-scalar returns False (NotImplemented).""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + assert (a == "hello") is False + assert (a != "hello") is True + assert (a == object()) is False + + def test_in_list_uses_equality(self): + """A matrix can be found in a list via value equality.""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + b = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + assert a in [Matrix(2, 2, [0.0, 0.0, 0.0, 0.0]), b] + + def test_unhashable(self): + """Defining value equality makes Matrix unhashable (mutable type).""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + with pytest.raises(TypeError): + hash(a) + with pytest.raises(TypeError): + _ = {a} + + def test_one_by_one_not_scalar_in_ordering(self): + """A 1x1 matrix only compares against another 1x1, not a larger shape.""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + one = Matrix(1, 1, [2.0]) + with pytest.raises(ValueError): + _ = a < one + # 1x1 vs 1x1 is fine. + assert Matrix(1, 1, [1.0]) < Matrix(1, 1, [2.0]) + + +class TestWhere: + """numpy-like Matrix.where(mask, a, b).""" + + def test_matrix_operands(self, shape, rng): + """where selects a where mask is non-zero and b elsewhere.""" + rows, cols = shape + mask_vals = [float(rng.randint(0, 1)) for _ in range(rows * cols)] + a_vals = [rng.uniform(-9, 9) for _ in range(rows * cols)] + b_vals = [rng.uniform(-9, 9) for _ in range(rows * cols)] + mask = Matrix(rows, cols, mask_vals) + a = Matrix(rows, cols, a_vals) + b = Matrix(rows, cols, b_vals) + result = Matrix.where(mask, a, b) + assert result.rows == rows + assert result.columns == cols + for i in range(rows): + for j in range(cols): + k = i * cols + j + expected = a_vals[k] if mask_vals[k] != 0.0 else b_vals[k] + assert result[i, j] == pytest.approx(expected) + + def test_scalar_a(self, shape, rng): + """where accepts a scalar for a, broadcast against the mask.""" + rows, cols = shape + mask_vals = [float(rng.randint(0, 1)) for _ in range(rows * cols)] + b_vals = [rng.uniform(-9, 9) for _ in range(rows * cols)] + mask = Matrix(rows, cols, mask_vals) + b = Matrix(rows, cols, b_vals) + result = Matrix.where(mask, 99.0, b) + for i in range(rows): + for j in range(cols): + k = i * cols + j + expected = 99.0 if mask_vals[k] != 0.0 else b_vals[k] + assert result[i, j] == pytest.approx(expected) + + def test_scalar_both(self): + """where accepts scalars for both a and b.""" + mask = Matrix(2, 2, [1.0, 0.0, 0.0, 1.0]) + result = Matrix.where(mask, 1.0, 0.0) + assert _flat(result) == [1.0, 0.0, 0.0, 1.0] + + def test_mask_from_comparison(self): + """A mask produced by a comparison method drives the selection.""" + a = Matrix(1, 4, [1.0, 5.0, 2.0, 8.0]) + result = Matrix.where(a.greater(3.0), a, 0.0) + assert _flat(result) == [0.0, 5.0, 0.0, 8.0] + + def test_nan_mask_selects_a(self): + """A NaN mask element is non-zero and selects a.""" + mask = Matrix(1, 2, [float("nan"), 0.0]) + result = Matrix.where(mask, 1.0, 2.0) + assert _flat(result) == [1.0, 2.0] + + def test_shape_mismatch_raises(self): + """A matrix operand whose shape differs from the mask raises.""" + mask = Matrix(2, 2, [1.0, 0.0, 0.0, 1.0]) + wrong = Matrix(1, 4, [1.0, 2.0, 3.0, 4.0]) + with pytest.raises(ValueError): + Matrix.where(mask, wrong, 0.0) + with pytest.raises(ValueError): + Matrix.where(mask, 0.0, wrong) + + def test_one_by_one_not_scalar(self): + """A 1x1 matrix operand must match the mask shape (not treated as scalar).""" + mask = Matrix(2, 2, [1.0, 0.0, 0.0, 1.0]) + with pytest.raises(ValueError): + Matrix.where(mask, Matrix(1, 1, [5.0]), 0.0) + + def test_bad_operand_type_raises(self): + """A non-matrix, non-scalar operand raises TypeError.""" + mask = Matrix(2, 2, [1.0, 0.0, 0.0, 1.0]) + with pytest.raises(TypeError): + Matrix.where(mask, "x", 0.0) + + +class TestComparisonMaskSelfBroadcast: + """Mask methods where self (the left operand) is the broadcast side. + + These pin the swap_right reflections inside Matrix_binary_op: when self is + the smaller operand the op is reflected before dispatch, so the result must + still read ``ref(self_elem, other_elem)``. + """ + + @pytest.mark.parametrize("name,ref", COMPARE_OPS, ids=[o[0] for o in COMPARE_OPS]) + def test_self_one_by_one(self, name, ref): + """A 1x1 self broadcasts against a full other.""" + full = Matrix(2, 3, [1.0, 5.0, 3.0, 4.0, 2.0, 6.0]) + s = Matrix(1, 1, [3.0]) + result = getattr(s, name)(full) + for i in range(2): + for j in range(3): + expected = 1.0 if ref(3.0, full[i, j]) else 0.0 + assert result[i, j] == expected + + @pytest.mark.parametrize("name,ref", COMPARE_OPS, ids=[o[0] for o in COMPARE_OPS]) + def test_self_row_vector(self, name, ref): + """A 1xN self broadcasts down the rows of an MxN other.""" + full = Matrix(2, 3, [1.0, 5.0, 3.0, 4.0, 2.0, 6.0]) + s = Matrix(1, 3, [3.0, 3.0, 3.0]) + result = getattr(s, name)(full) + for i in range(2): + for j in range(3): + expected = 1.0 if ref(s[0, j], full[i, j]) else 0.0 + assert result[i, j] == expected + + @pytest.mark.parametrize("name,ref", COMPARE_OPS, ids=[o[0] for o in COMPARE_OPS]) + def test_self_column_vector(self, name, ref): + """An Mx1 self broadcasts across the columns of an MxN other.""" + full = Matrix(2, 3, [1.0, 5.0, 3.0, 4.0, 2.0, 6.0]) + s = Matrix(2, 1, [3.0, 4.0]) + result = getattr(s, name)(full) + for i in range(2): + for j in range(3): + expected = 1.0 if ref(s[i, 0], full[i, j]) else 0.0 + assert result[i, j] == expected + + @pytest.mark.parametrize("name,ref", COMPARE_OPS, ids=[o[0] for o in COMPARE_OPS]) + def test_self_row_vs_column_outer(self, name, ref): + """A 1xN self against an Mx1 other forms the MxN outer-product mask.""" + s = Matrix(1, 3, [1.0, 3.0, 5.0]) + other = Matrix(2, 1, [2.0, 4.0]) + result = getattr(s, name)(other) + assert (result.rows, result.columns) == (2, 3) + for i in range(2): + for j in range(3): + expected = 1.0 if ref(s[0, j], other[i, 0]) else 0.0 + assert result[i, j] == expected + + +class TestComparisonMaskCoercion: + """Mask methods accept list/tuple and bool operands.""" + + def test_list_operand_row_broadcast(self): + """A list operand is taken as a row vector and broadcasts.""" + a = Matrix(2, 3, [1.0, 5.0, 3.0, 4.0, 2.0, 6.0]) + result = a.less([3.0, 3.0, 3.0]) + for i in range(2): + for j in range(3): + assert result[i, j] == (1.0 if a[i, j] < 3.0 else 0.0) + + def test_tuple_operand_same_shape(self): + """A tuple operand matching the flat shape compares element-wise.""" + a = Matrix(1, 3, [1.0, 3.0, 5.0]) + result = a.equal((1.0, 0.0, 5.0)) + assert _flat(result) == [1.0, 0.0, 1.0] + + def test_bool_operand_is_scalar(self): + """A bool operand is the scalar 1.0 (True) or 0.0 (False).""" + a = Matrix(1, 3, [0.0, 1.0, 2.0]) + assert _flat(a.less(True)) == [1.0, 0.0, 0.0] + assert _flat(a.greater(False)) == [0.0, 1.0, 1.0] + assert _flat(a.equal(True)) == [0.0, 1.0, 0.0] + + def test_empty_list_operand_raises(self): + """An empty list/tuple cannot be coerced and raises ValueError.""" + a = Matrix(1, 3, [0.0, 1.0, 2.0]) + with pytest.raises(ValueError): + a.less([]) + with pytest.raises(ValueError): + a.equal(()) + + def test_malformed_list_operand_raises(self): + """A list holding a non-number raises TypeError.""" + a = Matrix(1, 3, [0.0, 1.0, 2.0]) + with pytest.raises(TypeError): + a.greater([1.0, "x", 3.0]) + + +class TestLexicographicNaN: + """NaN handling in the lexicographic comparison operators.""" + + def test_all_nan_equals_itself(self): + """An all-NaN matrix compares == equal to itself (NaN never decides).""" + nan = float("nan") + a = Matrix(1, 3, [nan, nan, nan]) + b = Matrix(1, 3, [nan, nan, nan]) + assert a == b + assert not (a != b) + assert not (a < b) + assert not (a > b) + assert a <= b + assert a >= b + + def test_nan_skipped_first_real_difference_decides(self): + """A leading NaN is skipped; the first real ordering decides.""" + nan = float("nan") + a = Matrix(1, 3, [nan, 1.0, 9.0]) + b = Matrix(1, 3, [nan, 2.0, 0.0]) + assert a < b + assert not (a > b) + + def test_nan_scalar_path(self): + """The scalar comparison path also skips NaN elements.""" + nan = float("nan") + a = Matrix(1, 3, [nan, nan, nan]) + assert (a == nan) is True + assert (a != nan) is False + assert (a < nan) is False + assert (a > nan) is False + + def test_equal_mask_vs_eq_operator_nan_divergence(self): + """equal() sends NaN->0.0, but the == operator treats all-NaN as equal.""" + nan = float("nan") + a = Matrix(1, 2, [nan, nan]) + b = Matrix(1, 2, [nan, nan]) + assert _flat(a.equal(b)) == [0.0, 0.0] + assert (a == b) is True + + +class TestReflectedScalarCompare: + """Reflected scalar comparisons mirror the matrix-on-left form.""" + + @pytest.mark.parametrize("scalar", [0.0, 3.0, 5.0]) + def test_reflected_equivalence(self, scalar): + """``scalar OP a`` equals the reflected ``a OP' scalar`` per operator.""" + a = Matrix(1, 3, [1.0, 3.0, 5.0]) + assert (scalar < a) is (a > scalar) + assert (scalar <= a) is (a >= scalar) + assert (scalar > a) is (a < scalar) + assert (scalar >= a) is (a <= scalar) + assert (scalar == a) is (a == scalar) + assert (scalar != a) is (a != scalar) + + +class TestRichcompareCoercion: + """Comparison operators accept list/tuple operands and bool scalars.""" + + @pytest.mark.parametrize("name,op", RICHCOMPARE_OPS, ids=[o[0] for o in RICHCOMPARE_OPS]) + def test_list_operand_matches_matrix(self, name, op): + """A same-shape list operand compares like the equivalent matrix.""" + a = Matrix(1, 3, [1.0, 2.0, 3.0]) + seq = [1.0, 2.0, 4.0] + assert op(a, seq) is op(a, Matrix(1, 3, seq)) + + def test_tuple_operand(self): + """A tuple operand of matching shape compares lexicographically.""" + a = Matrix(1, 3, [1.0, 2.0, 3.0]) + assert (a == (1.0, 2.0, 3.0)) is True + assert (a < (1.0, 2.0, 4.0)) is True + + def test_list_shape_mismatch_eq_false(self): + """A list whose shape differs makes == False / != True (no raise).""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + assert (a == [1.0, 2.0, 3.0, 4.0]) is False + assert (a != [1.0, 2.0, 3.0, 4.0]) is True + + def test_list_shape_mismatch_ordering_raises(self): + """An ordering operator against a mismatched list raises ValueError.""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + with pytest.raises(ValueError): + _ = a < [1.0, 2.0, 3.0, 4.0] + + def test_bool_scalar(self): + """A bool operand is the scalar 1.0 (True) or 0.0 (False).""" + a = Matrix(1, 2, [1.0, 1.0]) + assert (a == True) is True # noqa: E712 — bool-as-scalar is intended + assert (a == False) is False # noqa: E712 — bool-as-scalar is intended + + def test_empty_list_eq_is_total(self): + """An empty list cannot be coerced, so == is False / != True (no raise).""" + a = Matrix(1, 3, [1.0, 2.0, 3.0]) + assert (a == []) is False + assert (a != []) is True + assert (a == ()) is False + assert (a != ()) is True + + def test_empty_list_ordering_raises(self): + """An ordering operator against an empty list raises ValueError.""" + a = Matrix(1, 3, [1.0, 2.0, 3.0]) + with pytest.raises(ValueError): + _ = a < [] + + def test_malformed_list_eq_is_total(self): + """A list holding a non-number makes == False / != True (no raise).""" + a = Matrix(1, 3, [1.0, 2.0, 3.0]) + assert (a == [1.0, "x", 3.0]) is False + assert (a != [1.0, "x", 3.0]) is True + + def test_nested_list_eq_is_total(self): + """A nested list (non-number elements) makes == False / != True.""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + assert (a == [[1.0, 2.0], [3.0, 4.0]]) is False + assert (a != [[1.0, 2.0], [3.0, 4.0]]) is True + + def test_malformed_list_ordering_raises(self): + """An ordering operator against a malformed list raises TypeError.""" + a = Matrix(1, 3, [1.0, 2.0, 3.0]) + with pytest.raises(TypeError): + _ = a < [1.0, "x", 3.0] + + def test_nested_list_ordering_raises(self): + """An ordering operator against a nested list raises TypeError.""" + a = Matrix(2, 2, [1.0, 2.0, 3.0, 4.0]) + with pytest.raises(TypeError): + _ = a < [[1.0, 2.0], [3.0, 4.0]] + + def test_in_list_with_uncoercible_elements(self): + """`matrix in container` stays exception-free past empty/malformed lists.""" + a = Matrix(1, 3, [1.0, 2.0, 3.0]) + b = Matrix(1, 3, [1.0, 2.0, 3.0]) + assert a in [[], [1.0, "x", 3.0], b] + + +class TestWhereExtended: + """where() NaN propagation, bool, and sequence operands.""" + + def test_nan_value_propagates(self): + """Selecting a NaN value propagates it; value NaN is not special-cased.""" + nan = float("nan") + mask = Matrix(1, 2, [0.0, 1.0]) + b = Matrix(1, 2, [nan, nan]) + result = Matrix.where(mask, 1.0, b) + assert math.isnan(result[0, 0]) + assert result[0, 1] == 1.0 + + def test_list_operand(self): + """A list value operand is taken as a row vector matching the mask.""" + mask = Matrix(1, 3, [1.0, 0.0, 1.0]) + result = Matrix.where(mask, [10.0, 20.0, 30.0], 0.0) + assert _flat(result) == [10.0, 0.0, 30.0] + + def test_tuple_operand(self): + """A tuple value operand is coerced the same as a list.""" + mask = Matrix(1, 3, [0.0, 1.0, 0.0]) + result = Matrix.where(mask, 9.0, (10.0, 20.0, 30.0)) + assert _flat(result) == [10.0, 9.0, 30.0] + + def test_list_shape_mismatch_raises(self): + """A list whose shape differs from the mask raises ValueError.""" + mask = Matrix(2, 2, [1.0, 0.0, 0.0, 1.0]) + with pytest.raises(ValueError): + Matrix.where(mask, [1.0, 2.0, 3.0], 0.0) + + def test_bool_operand(self): + """Bool value operands are scalars (True -> 1.0, False -> 0.0).""" + mask = Matrix(1, 2, [1.0, 0.0]) + result = Matrix.where(mask, True, False) + assert _flat(result) == [1.0, 0.0] + + def test_empty_list_operand_raises(self): + """An empty list value operand cannot be coerced and raises ValueError.""" + mask = Matrix(1, 3, [1.0, 0.0, 1.0]) + with pytest.raises(ValueError): + Matrix.where(mask, [], 0.0) + + def test_empty_list_b_operand_raises(self): + """An empty list in the b (second) position also raises ValueError.""" + mask = Matrix(1, 3, [1.0, 0.0, 1.0]) + with pytest.raises(ValueError): + Matrix.where(mask, 0.0, []) + + def test_malformed_list_operand_raises(self): + """A list value operand holding a non-number raises TypeError.""" + mask = Matrix(1, 3, [1.0, 0.0, 1.0]) + with pytest.raises(TypeError): + Matrix.where(mask, [1.0, "x", 3.0], 0.0) + + def test_malformed_list_b_operand_raises(self): + """A non-number element in the b (second) position raises TypeError.""" + mask = Matrix(1, 3, [1.0, 0.0, 1.0]) + with pytest.raises(TypeError): + Matrix.where(mask, 0.0, [1.0, "x", 3.0]) + + +class TestFmaBroadcast: + """fma() row- and column-vector broadcasting of b and c.""" + + def test_row_vector_b(self): + """A 1xN b broadcasts down the rows of self.""" + a = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + row = Matrix(1, 3, [10.0, 20.0, 30.0]) + out = a.fma(row, 0.0) + for i in range(2): + for j in range(3): + assert out[i, j] == ref_fma(a[i, j], row[0, j], 0.0) + + def test_column_vector_b(self): + """An Mx1 b broadcasts across the columns of self.""" + a = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + col = Matrix(2, 1, [10.0, 20.0]) + out = a.fma(col, 0.0) + for i in range(2): + for j in range(3): + assert out[i, j] == ref_fma(a[i, j], col[i, 0], 0.0) + + def test_row_vector_c(self): + """A 1xN c broadcasts down the rows of self.""" + a = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + row = Matrix(1, 3, [10.0, 20.0, 30.0]) + out = a.fma(1.0, row) + for i in range(2): + for j in range(3): + assert out[i, j] == ref_fma(a[i, j], 1.0, row[0, j]) + + def test_column_vector_c(self): + """An Mx1 c broadcasts across the columns of self.""" + a = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + col = Matrix(2, 1, [10.0, 20.0]) + out = a.fma(1.0, col) + for i in range(2): + for j in range(3): + assert out[i, j] == ref_fma(a[i, j], 1.0, col[i, 0]) + + def test_row_b_and_column_c(self): + """A row-vector b and a column-vector c broadcast independently.""" + a = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + row = Matrix(1, 3, [10.0, 20.0, 30.0]) + col = Matrix(2, 1, [100.0, 200.0]) + out = a.fma(row, col) + for i in range(2): + for j in range(3): + assert out[i, j] == ref_fma(a[i, j], row[0, j], col[i, 0]) + + def test_broadcast_in_place(self): + """A broadcast operand still supports in_place=True.""" + a = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + av = _flat(a) + row = Matrix(1, 3, [10.0, 20.0, 30.0]) + expected = [ref_fma(av[i * 3 + j], row[0, j], 0.0) + for i in range(2) for j in range(3)] + out = a.fma(row, 0.0, in_place=True) + assert out is a + assert _flat(a) == expected + + def test_broadcast_leaves_operand_unmodified(self): + """Materialising a broadcast operand does not mutate the source vector.""" + a = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + row = Matrix(1, 3, [10.0, 20.0, 30.0]) + before = _flat(row) + a.fma(row, 0.0) + assert _flat(row) == before + + def test_wrong_length_row_vector_raises(self): + """A 1xN b whose width mismatches self's columns raises ValueError.""" + a = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + bad = Matrix(1, 2, [1.0, 2.0]) + with pytest.raises(ValueError): + a.fma(bad, 0.0) + + def test_wrong_length_column_vector_raises(self): + """An Mx1 b whose height mismatches self's rows raises ValueError.""" + a = Matrix(2, 3, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + bad = Matrix(3, 1, [1.0, 2.0, 3.0]) + with pytest.raises(ValueError): + a.fma(bad, 0.0)