Add support for __matmul__ in ExpressionNode by alexfikl · Pull Request #235 · inducer/pymbolic

alexfikl · 2026-02-12T08:48:47Z

This adds some support for the @ operator to ExpressionNode with

A new Matmul node (naming? the other nodes don't use short names)
Support in some mappers for Matmul, mainly StringifyMapper and Collector and friends.

alexfikl · 2026-02-12T08:51:33Z

pymbolic/mapper/evaluator.py

+    @override
+    def map_matmul(self, expr: p.Matmul, /) -> ResultT:
+        if not expr.children:
+            return cast("ResultT", 1)


Not sure this is a good idea? Mainly because ndarray @ 1 fails due to the shape mismatch (even if it's (1,)).

Yeah, might be better to error in that case.

Removed in https://github.com/inducer/pymbolic/compare/b3b9070e6a5ed0d9ab1e65dde84a0a6a08432e74..7f9907d104dff38ce50756f92f994532e49a47ad.

alexfikl · 2026-02-12T08:54:13Z

pymbolic/mapper/flop_counter.py

-    map_product: Callable[[Self, p.Product], ArithmeticExpression] = map_sum
+    map_sum: Callable[[Self, p.Sum], ArithmeticExpression] = _count_children
+    map_product: Callable[[Self, p.Product], ArithmeticExpression] = _count_children
+    map_matmul: Callable[[Self, p.Matmul], ArithmeticExpression] = _count_children


Since this is a "matrix multiply", not sure this makes sense in the FlopCounter, but I guess subclasses will have to figure it out?

Yeah, leave unimplemented.

This was already implemented in CombineMapper, so all this did was add the nops + 1 to the children cost. I switched that to NotImplementedError completely to avoid confusion.

https://github.com/inducer/pymbolic/compare/b3b9070e6a5ed0d9ab1e65dde84a0a6a08432e74..7f9907d104dff38ce50756f92f994532e49a47ad

alexfikl · 2026-02-12T08:56:08Z

pymbolic/mapper/stringifier.py

+        enclosing_prec: int, *args: P.args, **kwargs: P.kwargs
+    ) -> str:
+        return self.parenthesize_if_needed(
+            self.join_rec(" ", expr.children, PREC_PRODUCT, *args, **kwargs),


? Not sure I have any better idea of what to map it to in LaTeX.

I think that's OK.

Added a \cdot like suggested below.

https://github.com/inducer/pymbolic/compare/b3b9070e6a5ed0d9ab1e65dde84a0a6a08432e74..7f9907d104dff38ce50756f92f994532e49a47ad

Copilot

Pull request overview

Adds a new Matmul expression primitive and plumbing so ExpressionNode can represent/use Python’s @ operator across core mappers (stringification, evaluation, traversal, etc.).

Changes:

Introduces prim.Matmul and implements ExpressionNode.__matmul__ / __rmatmul__.
Adds Matmul handling across multiple mappers (stringifiers, evaluator, graphviz, flop counter, generic mapper base classes).
Adds a basic unit test and updates the basedpyright baseline.

Reviewed changes

Copilot reviewed 17 out of 17 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
`pymbolic/primitives.py`	Adds `Matmul` node and `@` operator overloads on `ExpressionNode`.
`pymbolic/mapper/__init__.py`	Adds `map_matmul` to base/utility mappers (default, combine, identity, walk, constant).
`pymbolic/mapper/stringifier.py`	Adds `map_matmul` for standard and LaTeX stringification.
`pymbolic/mapper/evaluator.py`	Adds evaluation logic for `Matmul` using `operator.matmul`/`reduce`; adjusts product evaluation implementation.
`pymbolic/mapper/graphviz.py`	Adds graphviz visualization for `Matmul`.
`pymbolic/mapper/flop_counter.py`	Counts `Matmul` similarly to `Sum`/`Product` (n-1 plus children).
`pymbolic/mapper/c_code.py`	Explicitly rejects `Matmul` in C codegen via `NotImplementedError`.
`pymbolic/mapper/cse_tagger.py`	Tags `Matmul` subexpressions for CSE handling.
`pymbolic/mapper/unifier.py`	Refactors imports to use `pymbolic.primitives as p` consistently.
`pymbolic/mapper/substitutor.py`	Minor import refactor for `Variable`.
`pymbolic/mapper/flattener.py`	Refactors to call `p.flattened_sum/product` directly.
`pymbolic/mapper/distributor.py`	Refactors `Sum` reference to `p.Sum`.
`pymbolic/mapper/constant_folder.py`	Refactors to use `prim.flattened_sum/product` directly.
`pymbolic/mapper/collector.py`	Import/type-expression refactors in `TermCollector.split_term`.
`pymbolic/mapper/coefficient.py`	Refactors to use `p.*` helpers/types directly.
`test/test_pymbolic.py`	Adds a basic `test_matmul` asserting node type and `str()` output.
`.basedpyright/baseline.json`	Adds new suppressions (unannotated class attribute / `Any`).

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

.basedpyright/baseline.json

pymbolic/mapper/stringifier.py

Copilot · 2026-02-12T15:39:56Z

pymbolic/mapper/stringifier.py

+        self, expr: p.Matmul, /, enclosing_prec: int, *args: P.args, **kwargs: P.kwargs
+    ) -> str:
+        return self.parenthesize_if_needed(
+            self.join_rec(" @ ", expr.children, PREC_PRODUCT, *args, **kwargs),


StringifyMapper.map_product uses join_rec_with_parens_around_types to add parentheses around quotient-like children even though they share precedence, improving readability and avoiding ambiguous a / b * c-style output. map_matmul currently uses join_rec directly, so expressions like (a/b) @ c may stringify as a / b @ c without clarifying grouping. Consider using the same parenthesization helper/operand-type handling as map_product for consistency.

Suggested change

self.join_rec(" @ ", expr.children, PREC_PRODUCT, *args, **kwargs),

self.join_rec_with_parens_around_types(

" @ ",

expr.children,

PREC_PRODUCT,

self.multiplicative_primitives,

*args,

**kwargs,

),

Done in https://github.com/inducer/pymbolic/compare/b3b9070e6a5ed0d9ab1e65dde84a0a6a08432e74..7f9907d104dff38ce50756f92f994532e49a47ad.

pymbolic/mapper/evaluator.py

test/test_pymbolic.py

.basedpyright/baseline.json

inducer · 2026-02-12T17:39:47Z

Thx!

feat: add support for __matmul__ in ExpressionNode

e2aadc3

alexfikl commented Feb 12, 2026

View reviewed changes

inducer requested a review from Copilot February 12, 2026 15:33

Copilot started reviewing on behalf of inducer February 12, 2026 15:34 View session

Copilot AI reviewed Feb 12, 2026

View reviewed changes

alexfikl added 4 commits February 12, 2026 19:19

feat: add support for new Matmul node in mappers

96a81da

feat: remove some unneeded imports

8c22ba0

test: add simple test for Matmul node

9dd86ca

chore: update baseline

7f9907d

alexfikl force-pushed the add-matmul branch from b3b9070 to 7f9907d Compare February 12, 2026 17:20

inducer enabled auto-merge (rebase) February 12, 2026 17:39

inducer merged commit c5483ac into inducer:main Feb 12, 2026
10 checks passed

alexfikl deleted the add-matmul branch February 12, 2026 18:39

Comments

Conversation

alexfikl commented Feb 12, 2026

Uh oh!

alexfikl Feb 12, 2026 • 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!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI Feb 12, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

inducer commented Feb 12, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

alexfikl Feb 12, 2026 •

edited

Loading