From dc918fe7a7fca5ee980c5360db1e5ea05c3bc155 Mon Sep 17 00:00:00 2001 From: Noureldin Date: Fri, 9 Jan 2026 11:04:37 -0800 Subject: [PATCH 1/3] return NotImplemted for unsupported operations in __array_ufunc__ --- tunits/core/cython/with_unit_value_array.pyx | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tunits/core/cython/with_unit_value_array.pyx b/tunits/core/cython/with_unit_value_array.pyx index 2e7fcb1..8d89225 100644 --- a/tunits/core/cython/with_unit_value_array.pyx +++ b/tunits/core/cython/with_unit_value_array.pyx @@ -133,7 +133,7 @@ class ValueArray(WithUnit): if self._is_dimensionless(): return getattr(ufunc, method)(*(np.asarray(x) for x in inputs), **kwargs) - raise NotImplemented + return NotImplemented @property def dtype(WithUnit self) -> np.dtype: @@ -160,4 +160,4 @@ class ValueArray(WithUnit): def to_proto(self, msg: Optional['tunits_pb2.ValueArray'] = None) -> 'tunits_pb2.ValueArray': ret = _ndarray_to_proto(self.value, msg) ret.units.extend(_units_to_proto(self.display_units)) - return ret \ No newline at end of file + return ret From d1ae49b230c09d3158ed7f4d3102c06a163b3938 Mon Sep 17 00:00:00 2001 From: Nour Yosri Date: Fri, 9 Jan 2026 12:06:36 -0800 Subject: [PATCH 2/3] implement matmul --- test/test_value_array.py | 7 +++++++ tunits/core/cython/with_unit_value_array.pyx | 11 +++++++++++ 2 files changed, 18 insertions(+) diff --git a/test/test_value_array.py b/test/test_value_array.py index a69708e..be10b0b 100644 --- a/test/test_value_array.py +++ b/test/test_value_array.py @@ -293,3 +293,10 @@ def test_unique() -> None: unit = tu.MHz v_arr = xs * unit assert np.array_equal(v_arr.unique(), np.unique(xs) * unit) + + +def test_matmul() -> None: + a = np.random.random((3, 4)) + b = np.random.random((4, 3)) * tu.ns + assert (a @ b).allclose((a @ b[tu.us]) * tu.us) + assert (b @ a).allclose((b[tu.s] @ a) * tu.s) diff --git a/tunits/core/cython/with_unit_value_array.pyx b/tunits/core/cython/with_unit_value_array.pyx index 8d89225..c97fb0f 100644 --- a/tunits/core/cython/with_unit_value_array.pyx +++ b/tunits/core/cython/with_unit_value_array.pyx @@ -116,6 +116,10 @@ class ValueArray(WithUnit): return self ** 2 if ufunc == np.reciprocal: return self.__rtruediv__(1) + if ufunc == np.matmul: + if isinstance(inputs[0], ValueArray): + return inputs[0].__matmul__(inputs[1]) + return inputs[1].__rmatmul__(inputs[0]) if ufunc in [ np.greater, @@ -161,3 +165,10 @@ class ValueArray(WithUnit): ret = _ndarray_to_proto(self.value, msg) ret.units.extend(_units_to_proto(self.display_units)) return ret + + def __matmul__(WithUnit self, other: np.ndarray): + return self.__with_value(self.value @ other) + + + def __rmatmul__(WithUnit self, other: np.ndarray): + return self.__with_value(other @ self.value) From 080f7cae988dcfc7b863e0a76502662b4a28184c Mon Sep 17 00:00:00 2001 From: Nour Yosri Date: Fri, 9 Jan 2026 12:12:43 -0800 Subject: [PATCH 3/3] nit --- test/test_value_array.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/test/test_value_array.py b/test/test_value_array.py index be10b0b..390e2cb 100644 --- a/test/test_value_array.py +++ b/test/test_value_array.py @@ -298,5 +298,7 @@ def test_unique() -> None: def test_matmul() -> None: a = np.random.random((3, 4)) b = np.random.random((4, 3)) * tu.ns - assert (a @ b).allclose((a @ b[tu.us]) * tu.us) - assert (b @ a).allclose((b[tu.s] @ a) * tu.s) + c: tu.TimeArray = a @ b # type: ignore[assignment] + d: tu.TimeArray = b @ a # type: ignore[assignment] + assert c.allclose((a @ b[tu.us]) * tu.us) + assert d.allclose((b[tu.s] @ a) * tu.s)