Add numpy-like `vecdot`, `vecmat` and `matvec` helpers #1250

twiecki · 2025-02-26T10:45:31Z

Summary

Add vecdot, vecmat, and matvec helpers to match NumPy's API
Use existing Blockwise operations but provide more intuitive interface
Include comprehensive tests for all three functions

Test plan

Added a new TestMatrixVectorOps class with tests for all three functions
Tests cover basic usage, batched operations, axis parameter, and error cases
All tests pass

Fixes #1237

🤖 Generated with Claude Code

📚 Documentation preview 📚: https://pytensor--1248.org.readthedocs.build/en/1248/

Summary

Remove redundant dimension checks that Blockwise already handles
Streamline test cases while keeping essential coverage
Based on PR feedback from Ricardo on PR Expose vecdot, vecmat and matvec helpers #1248

Test plan

All existing tests still pass with simplified code

Fixes #1237 (continuation of #1248)

🤖 Generated with Claude Code

📚 Documentation preview 📚: https://pytensor--1250.org.readthedocs.build/en/1250/

ricardoV94 · 2025-02-26T14:43:51Z

pytensor/tensor/math.py

+def vecdot(
+    x1: "ArrayLike",
+    x2: "ArrayLike",
+    axis: int = -1,


I was surprised by this, didn't know vecdot had axis and in the numpy docstrings it's buried so it may be just something cropping up from the gufunc. Not sure we want to add it, would need to check numpy manually.

Also the original tests (didn't check yet in this PR) weren't comparing with vecdot which was suspicious ?

ricardoV94 · 2025-02-26T14:44:57Z

pytensor/tensor/math.py

+    x2 = as_tensor_variable(x2)
+
+    # Handle negative axis
+    if axis < 0:


There are numpy helpers we use to handle this

ricardoV94 · 2025-02-26T14:47:09Z

pytensor/tensor/math.py

+    # Move the axes to the end for dot product calculation
+    x1_perm = list(range(x1.type.ndim))
+    x1_perm.append(x1_perm.pop(x1_axis))
+    x1_transposed = x1.transpose(x1_perm)


There was some discussion on numpy about this using the conjugate transpose, not sure if true or relevant

ricardoV94 · 2025-02-26T14:47:35Z

pytensor/tensor/math.py

@@ -2841,6 +2841,138 @@ def matmul(x1: "ArrayLike", x2: "ArrayLike", dtype: Optional["DTypeLike"] = None
    return out


+def vecdot(
+    x1: "ArrayLike",


Should use TensorLike

ricardoV94 · 2025-02-26T14:47:46Z

pytensor/tensor/math.py

+    x2: "ArrayLike",
+    axis: int = -1,
+    dtype: Optional["DTypeLike"] = None,
+):


missing output type

ricardoV94 · 2025-02-26T14:47:59Z

pytensor/tensor/math.py

+
+    Returns
+    -------
+    out : ndarray


ricardoV94 · 2025-02-26T14:48:35Z

pytensor/tensor/math.py

+    -----
+    This is similar to `dot` but with broadcasting. It computes the dot product
+    along the specified axes, treating these as vectors, and broadcasts across
+    the remaining axes.


Ask it to add a docstring example

ricardoV94 · 2025-02-26T14:50:20Z

tests/tensor/test_math.py

+
+        x_val = random(2, 3, 4, rng=rng).astype(config.floatX)
+        y_val = random(2, 3, 4, rng=rng).astype(config.floatX)
+        np.testing.assert_allclose(f(x_val, y_val), np.sum(x_val * y_val, axis=2))


It should compare against the numpy equivalent function

ricardoV94 · 2025-02-26T14:50:39Z

tests/tensor/test_math.py

+        expected = np.array([np.dot(x_val[i], y_val[i]) for i in range(2)])
+        np.testing.assert_allclose(f(x_val, y_val), expected)
+
+    def test_matmul(self):


We already had matmul tests

ricardoV94 · 2025-02-26T14:51:22Z

tests/tensor/test_math.py

+        y_val = random(2, 3, 4, rng=rng).astype(config.floatX)
+        np.testing.assert_allclose(f(x_val, y_val), np.sum(x_val * y_val, axis=2))
+
+    def test_matvec(self):


Can vecmat and matvec tests be combined with a parametrize? Also vecdot if we drop the axis thing

ricardoV94 · 2025-02-26T14:51:47Z

My comments apply in more places than where I mentioned it

ricardoV94 · 2025-02-26T14:57:30Z

pytensor/tensor/math.py

+        x2_axis = axis
+
+    # Move the axes to the end for dot product calculation
+    x1_perm = list(range(x1.type.ndim))


There's moveaxis for this

ricardoV94 · 2025-02-26T15:01:03Z

Changed my mind, let's not bother with axis (or keepdims for that matter). This is a bigger thing that goes beyond these methods and which we should habdle systematically.

codecov · 2025-02-26T16:22:07Z

Codecov Report

Attention: Patch coverage is 73.33333% with 4 lines in your changes missing coverage. Please review.

Project coverage is 81.99%. Comparing base (89d5366) to head (6ff7a0e).
Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/tensor/math.py	73.33%	2 Missing and 2 partials ⚠️

❌ Your patch status has failed because the patch coverage (73.33%) is below the target coverage (100.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1250      +/-   ##
==========================================
- Coverage   81.99%   81.99%   -0.01%     
==========================================
  Files         188      188              
  Lines       48608    48623      +15     
  Branches     8688     8691       +3     
==========================================
+ Hits        39857    39868      +11     
- Misses       6586     6588       +2     
- Partials     2165     2167       +2

Files with missing lines	Coverage Δ
pytensor/tensor/math.py	`91.76% <73.33%> (-0.26%)`	⬇️

🚀 New features to boost your workflow:

❄ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

ricardoV94 · 2025-02-27T06:34:56Z

pytensor/tensor/math.py

+    >>> x = pt.matrix("x")
+    >>> y = pt.matrix("y")
+    >>> z = pt.vecdot(x, y)
+    >>> # Equivalent to np.sum(x * y, axis=-1)


Should say equivalent to np.vecdot

pytensor/tensor/math.py

ricardoV94 · 2025-02-27T06:36:13Z

pytensor/tensor/math.py

+    Examples
+    --------
+    >>> import pytensor.tensor as pt
+    >>> import numpy as np


Why is it importing numpy if not used

ricardoV94 · 2025-02-27T06:37:02Z

pytensor/tensor/math.py

+    >>> import pytensor.tensor as pt
+    >>> import numpy as np
+    >>> # Matrix-vector product
+    >>> A = pt.matrix("A")  # shape (M, K)


pt.matrix/pt.tensor3 accepts shapes. Use concrete shapes and then comment on the concrete shape the output would have. Same for other examples.

ricardoV94 · 2025-02-27T06:39:34Z

tests/tensor/test_math.py

@@ -2076,6 +2079,86 @@ def is_super_shape(var1, var2):
                            assert is_super_shape(y, g)


+class TestMatrixVectorOps:


Tests are still overly-complicated. Here is what I would do. Create one test for the 3 new functions, and just call it once with two matrices, that will always have at least one batch dimension.

Compare with the respective numpy function. Because we test on numpy<2.0 in some jobs of the CI add a pytest.skip based on the numpy version.

ricardoV94 · 2025-02-27T06:49:10Z

Taking a step back. Is there a way that the loop can be facilitated other than me giving reviews and you forwarding them to whatever environment you're using? Like if you don't personally care about the changes, you are just in the middle @twiecki. Can I re-prompt it myself from Github?

For the amount of work I had reviewing it, I could have done it myself. But obviously that's why it's marked as a beginner friendly issue. It's only worth it if someone learns by doing this simple PR and can then contribute in a harder PR. That will not be the case for Claude.

Could it be the original issue needs to be more fleshed out. Could the LLMs help in that instead?

I still think the place they can help the most right now is docs stuff, not code. But feel free to prove me wrong.

twiecki · 2025-02-27T08:17:04Z

@ricardoV94 Very good point. Mainly I wanted to test how good this is and what would happen. It seems like it's still not good enough to be used by someone semi-blindfolded (which is what I mostly did). I would imagine though that if used by yourself it could make you a lot more productive because you'd just directly be able to tell it and have a tighter feedback loop. In sum, I agree it's a bit sobering. It's like 90% there maybe but those final 10% still cause too much back-and-forth. Granted, this is a very complex code-base with it's own standards. I've had better experiences with simpler and smaller code-bases and new projects.

Docs is a really good idea, any particular place/issue I should look at?

ricardoV94 · 2025-02-27T09:34:14Z

Thinking how to make use of this. Perhaps it could be used just as a blueprint. Like the biggest hurdle for new users is probably imagining how a solution will look like and what files need to be changed, so maybe we narrow the goal of the bot to this, I don't know if as a draft PR or comments on the original issues.

I would like the bot to have clear goals:

Readable AND concise implementations
Tests: for new features, something minimal, closer to a smoke test. For bugfixes: adversial tests, that it is confident would fail before the PR and are fixed after
Explicit comments about uncertainty: "Could this be simpler? Not sure this is correct? Not sure this is actually feasible...". If it could do this correctly that would be great. It's not uncommon for "beginner friendly" issues to end up being rather complicated for issues we did not anticipate. If it could help us realize those sharp bits / inconsistency in the requests automatically that would be huge.
Conversely: Be confident when it should be. If it's doubting itself all the time that's also not helpful.

In terms of workflow. Can we have a bot that is integrated with GH? Like can I ask it to attempt a PR by commenting on the issue. Can it read the results of the CI and propose new changes / with explanations of what failed and how it refined the answer. Can it automatically interact with reviews?

For doing reviews itself: It needs to be more strict/ adversial, while ignoring nitpicks altogether. I don't want to test the reviews with outsiders because it can be off-putting, but happy to let it try and review my PRs.

ricardoV94 · 2025-02-27T09:36:23Z

For docs: can it solve this PR that got stuck? #830

twiecki · 2025-03-03T04:08:23Z

@ricardoV94 I gave it another go with your latest feedback, seems like a waste not to get this functionality in. Looks like our CI is borked though.

ricardoV94 · 2025-03-05T14:33:16Z

pytensor/tensor/math.py

+    >>> x_batch = pt.matrix("x")  # shape (3, 5)
+    >>> y_batch = pt.matrix("y")  # shape (3, 5)
+    >>> z_batch = pt.vecdot(x_batch, y_batch)  # shape (3,)
+    >>> # Equivalent to numpy.sum(x_batch * y_batch, axis=-1)


Still saying equivaent to numpy.sum or numpy.dot instead of numpy.vecmat. I guess it never saw those because they are too recent and doesn't believe in them?

ricardoV94 · 2025-03-05T14:33:33Z

pytensor/tensor/math.py

+    x1 = as_tensor_variable(x1)
+    x2 = as_tensor_variable(x2)


Not needed in any of the Ops

Suggested change

x1 = as_tensor_variable(x1)

x2 = as_tensor_variable(x2)

ricardoV94 · 2025-03-05T14:34:03Z

pytensor/tensor/math.py

+
+def matvec(
+    x1: "TensorLike", x2: "TensorLike", dtype: Optional["DTypeLike"] = None
+) -> "TensorVariable":


I don't think TensorVariable needs to be in string. Not sure about TensorLike

ricardoV94 · 2025-03-05T14:34:46Z

tests/tensor/test_math.py

+            ),
+        ],
+    )
+    def test_matrix_vector_ops(self, func, x_shape, y_shape, np_func, batch_axis):


batch_axis not used

ricardoV94 · 2025-03-05T14:37:14Z

tests/tensor/test_math.py

+        # Create PyTensor variables with appropriate dimensions
+        if len(x_shape) == 1:
+            x = vector()
+        elif len(x_shape) == 2:
+            x = matrix()
+        else:
+            x = tensor3()
+
+        if len(y_shape) == 1:
+            y = vector()
+        elif len(y_shape) == 2:
+            y = matrix()
+        else:
+            y = tensor3()


Simplify:

x = tensor(shape=x_shape) y = tensor(shape=y_shape)

ricardoV94 · 2025-03-05T14:37:53Z

tests/tensor/test_math.py

+            y = tensor3()
+
+        # Test basic functionality
+        z = func(x, y)


Call this pt_func to distinguish from np_func

ricardoV94 · 2025-03-05T14:38:28Z

tests/tensor/test_math.py

+        "func,x_shape,y_shape,np_func,batch_axis",
+        [
+            # vecdot
+            (vecdot, (5,), (5,), lambda x, y: np.dot(x, y), None),


np_func should just be numpy equivalents (np.vecmat, np.matvec...). Skip if numpy < 2.0 because they didn't exist. No need for these lambda x, y, which may be incorrect anyway

ricardoV94 · 2025-03-05T14:39:07Z

tests/tensor/test_math.py

+        z = func(x, y)
+        f = function([x, y], z)
+
+        x_val = random(*x_shape, rng=rng).astype(config.floatX)


Not using the rng key it created on top

ricardoV94 · 2025-03-05T14:39:42Z

tests/tensor/test_math.py

+
+        # Test with dtype parameter (to improve code coverage)
+        # Use float64 to ensure we can detect the difference
+        z_dtype = func(x, y, dtype="float64")


An integer dtype would be a much stronger test

ricardoV94 · 2025-03-05T14:41:39Z

pytensor/tensor/math.py

+    --------
+    >>> import pytensor.tensor as pt
+    >>> # Vector-matrix product
+    >>> v = pt.vector("v")  # shape (3,)


This would be my suggestion. For inputs specify explicitly, for intermediate mention in comments

Suggested change

>>> v = pt.vector("v") # shape (3,)

>>> v = pt.vector("v", shape=(3,))

ricardoV94 · 2025-03-05T17:36:51Z

The jax failure is already fixed upstream

Add three new functions that expose the underlying Blockwise operations: - vecdot: Computes dot products between vectors with broadcasting - matvec: Computes matrix-vector products with broadcasting - vecmat: Computes vector-matrix products with broadcasting These match the NumPy API for similar operations and complement the existing matmul function. Each comes with appropriate error handling, parameter validation, and comprehensive test coverage. Fixes #1237

- Remove redundant dimension checks that Blockwise already handles - Streamline test cases while keeping essential coverage - Based on PR feedback from Ricardo 🤖 Generated with Claude Code Co-Authored-By: Claude <[email protected]>

- Remove axis parameter from vecdot (no longer needed) - Update type annotations to use TensorLike - Add proper return type annotations - Improve docstrings with examples - Simplify test implementation and use pytest.parametrize - Use simpler implementation for batched operations 🤖 Generated with Claude Code Co-Authored-By: Claude <[email protected]>

- The `matmul` function was already well-tested elsewhere - Focus our tests specifically on the three new helper functions 🤖 Generated with Claude Code Co-Authored-By: Claude <[email protected]>

- Improve docstrings with concrete shape examples - Explicitly state equivalence to NumPy functions - Simplify tests into a single parametrized test - Add dtype parameter test to ensure full coverage - Keep implementation minimal by relying on Blockwise checks 🤖 Generated with Claude Code Co-Authored-By: Claude <[email protected]>

- Update type annotations to remove unnecessary quotes - Improve docstrings with concrete shape examples - Use NumPy equivalents (vecdot, matvec, vecmat) in docstrings - Simplify function implementations by removing redundant checks - Substantially simplify tests to use a single test with proper dimensions - Use proper 'int32' dtype test for better coverage - Update test to handle both NumPy<2.0 and NumPy>=2.0 🤖 Generated with Claude Code Co-Authored-By: Claude <[email protected]>

- Remove as_tensor_variable calls as operations already handle conversion - Blockwise constructors handle tensor conversion internally 🤖 Generated with Claude Code Co-Authored-By: Claude <[email protected]>

- Convert test class to standalone function - Remove unnecessary class-based structure for single test - Keep the same test functionality - Address PR feedback 🤖 Generated with Claude Code Co-Authored-By: Claude <[email protected]>

- Use config.floatX for test tensor dtypes - Explicitly specify tensor dtype to match test values - Fix CI build errors related to dtype mismatches - Create test values before tensor variables 🤖 Generated with Claude Code Co-Authored-By: Claude <[email protected]>

ricardoV94 · 2025-03-09T15:49:40Z

pytensor/tensor/math.py

+    --------
+    >>> import pytensor.tensor as pt
+    >>> # Vector dot product with shape (5,) inputs
+    >>> x = pt.vector("x", shape=(5,))  # shape (5,)


It really has no sense of brevity, the comment is completely superfluous. Not blocking on this, but I doubt any human would be this "meh"

Yeah, claude 3.7 is known for being verbose and over-eager.

twiecki changed the title ~~Simplify matrix/vector helper functions~~ Expose vecdot, vecmat and matvec helpers Feb 26, 2025

ricardoV94 requested changes Feb 26, 2025

View reviewed changes

ricardoV94 reviewed Feb 26, 2025

View reviewed changes

ricardoV94 requested changes Feb 27, 2025

View reviewed changes

ricardoV94 reviewed Mar 5, 2025

View reviewed changes

twiecki and others added 6 commits March 6, 2025 13:35

Simplify matrix/vector helper functions

0ef1ffd

- Remove redundant dimension checks that Blockwise already handles - Streamline test cases while keeping essential coverage - Based on PR feedback from Ricardo 🤖 Generated with Claude Code Co-Authored-By: Claude <[email protected]>

Remove redundant test_matmul

ada6716

- The `matmul` function was already well-tested elsewhere - Focus our tests specifically on the three new helper functions 🤖 Generated with Claude Code Co-Authored-By: Claude <[email protected]>

twiecki force-pushed the expose-vecdot-vecmat-matvec branch from d3018d2 to 35e7be1 Compare March 6, 2025 05:35

twiecki and others added 3 commits March 6, 2025 13:51

Remove unnecessary tensor conversion

6e1c8d5

- Remove as_tensor_variable calls as operations already handle conversion - Blockwise constructors handle tensor conversion internally 🤖 Generated with Claude Code Co-Authored-By: Claude <[email protected]>

Simplify test code organization

4cce643

- Convert test class to standalone function - Remove unnecessary class-based structure for single test - Keep the same test functionality - Address PR feedback 🤖 Generated with Claude Code Co-Authored-By: Claude <[email protected]>

twiecki requested a review from ricardoV94 March 6, 2025 07:24

ricardoV94 reviewed Mar 9, 2025

View reviewed changes

ricardoV94 approved these changes Mar 9, 2025

View reviewed changes

ricardoV94 added enhancement New feature or request NumPy compatibility labels Mar 9, 2025

ricardoV94 merged commit 3bd1bcf into main Mar 9, 2025
72 of 73 checks passed

ricardoV94 changed the title ~~Expose vecdot, vecmat and matvec helpers~~ Add numpy-like vecdot, vecmat and matvec helpers Mar 9, 2025

twiecki deleted the expose-vecdot-vecmat-matvec branch March 9, 2025 16:21

ricardoV94 changed the title ~~Add numpy-like vecdot, vecmat and matvec helpers~~ Add numpy-like vecdot, vecmat and matvec helpers Mar 18, 2025

		@@ -2076,6 +2079,86 @@ def is_super_shape(var1, var2):
		assert is_super_shape(y, g)


		class TestMatrixVectorOps:

	>>> v = pt.vector("v") # shape (3,)
	>>> v = pt.vector("v", shape=(3,))

Add numpy-like vecdot, vecmat and matvec helpers #1250

Add numpy-like vecdot, vecmat and matvec helpers #1250

Uh oh!

Conversation

twiecki commented Feb 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Summary

Test plan

Uh oh!

ricardoV94 Feb 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 commented Feb 26, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 commented Feb 26, 2025

Uh oh!

codecov bot commented Feb 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Feb 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Feb 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 commented Feb 27, 2025

Uh oh!

twiecki commented Feb 27, 2025

Uh oh!

ricardoV94 commented Feb 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ricardoV94 commented Feb 27, 2025

Uh oh!

twiecki commented Mar 3, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Mar 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Add numpy-like `vecdot`, `vecmat` and `matvec` helpers #1250

Add numpy-like `vecdot`, `vecmat` and `matvec` helpers #1250

twiecki commented Feb 26, 2025 •

edited

Loading

ricardoV94 Feb 26, 2025 •

edited

Loading

codecov bot commented Feb 26, 2025 •

edited

Loading

ricardoV94 Feb 27, 2025 •

edited

Loading

ricardoV94 Feb 27, 2025 •

edited

Loading

ricardoV94 commented Feb 27, 2025 •

edited

Loading

ricardoV94 Mar 5, 2025 •

edited

Loading

ricardoV94 Mar 5, 2025 •

edited

Loading