Implement `aesara.tensor.matmul` #744

zoj613 · 2022-01-11T22:29:17Z

closes #488

This implements an aesara equivalent of np.matmul.

The behavior depends on the arguments in the following way.

If both arguments are 2-D they are multiplied like conventional matrices.

If either argument is N-D, N > 2, it is treated as a stack of matrices residing in the last two indexes and broadcast accordingly.

If the first argument is 1-D, it is promoted to a matrix by prepending a 1 to its dimensions. After matrix multiplication the prepended 1 is removed.

If the second argument is 1-D, it is promoted to a matrix by appending a 1 to its dimensions. After matrix multiplication the appended 1 is removed.

matmul differs from dot in two important ways:

Multiplication by scalars is not allowed.

Stacks of matrices are broadcast together as if the matrices were elements, respecting the signature (n,k),(k,m)->(n,m)

References

https://numpy.org/doc/stable/reference/generated/numpy.matmul.html

codecov · 2022-01-11T23:24:07Z

Codecov Report

Merging #744 (5abde4d) into main (8763981) will increase coverage by 0.01%.
The diff coverage is 100.00%.

❗ Current head 5abde4d differs from pull request most recent head b60030f. Consider uploading reports for the commit b60030f to get more accurate results

@@            Coverage Diff             @@
##             main     #744      +/-   ##
==========================================
+ Coverage   79.23%   79.25%   +0.01%     
==========================================
  Files         152      152              
  Lines       47943    48006      +63     
  Branches    10909    10933      +24     
==========================================
+ Hits        37990    38048      +58     
+ Misses       7453     7449       -4     
- Partials     2500     2509       +9

Impacted Files	Coverage Δ
aesara/tensor/nlinalg.py	`98.76% <100.00%> (+0.18%)`	⬆️
aesara/sparse/type.py	`72.11% <0.00%> (-2.66%)`	⬇️
aesara/link/numba/dispatch/tensor_basic.py	`97.95% <0.00%> (-2.05%)`	⬇️
aesara/link/basic.py	`85.00% <0.00%> (-1.50%)`	⬇️
aesara/tensor/shape.py	`90.93% <0.00%> (-1.29%)`	⬇️
aesara/tensor/subtensor_opt.py	`86.32% <0.00%> (-0.80%)`	⬇️
aesara/link/c/lazylinker_c.py	`65.95% <0.00%> (-0.71%)`	⬇️
aesara/link/c/cutils.py	`68.18% <0.00%> (-0.71%)`	⬇️
aesara/sparse/basic.py	`82.47% <0.00%> (-0.43%)`	⬇️
aesara/printing.py	`49.52% <0.00%> (-0.24%)`	⬇️
... and 27 more

brandonwillard

Looks great!

The next two important methods that need to be implemented are Op.grad (or Op.L_op) and Op.infer_shape.

They're both optional, but they can really make or break the usefulness of an Op. At the very least, it's good to have "stubs" for those that explicitly say they aren't implemented.

You can use tests.unittest_tools.verify_grad to create numeric gradient tests (or even tests.tensor.utils.makeTester for a basic test suite), and tests.unittest_tools.InferShapeTester to create automate the Op.infer_shape tests.

brandonwillard

Is it possible that np.matmul functionality can be implemented by a helper function that dispatches to the existing aesara.tensor.math.Dot and aesara.tensor.math.tensordot? If so, that would save considerable time and effort attempting to (re)implement the Op.grad and Op.infer_shape logic.

aesara/tensor/__init__.py

aesara/tensor/nlinalg.py

tests/tensor/test_nlinalg.py

zoj613 · 2022-01-12T21:43:47Z

Is it possible that np.matmul functionality can be implemented by a helper function that dispatches to the existing aesara.tensor.math.Dot and aesara.tensor.math.tensordot? If so, that would save considerable time and effort attempting to (re)implement the Op.grad and Op.infer_shape logic.

Not sure this is best given that matmul behaviour is different from dot's depending on the shape of the inputs. See the PR description. I managed to implement the infer_shape method without much hassle.

zoj613 · 2022-01-12T21:53:40Z

tests.unittest_tools.InferShapeTester to create automate the Op.infer_shape tests.

Could you elaborate on this part? Do I necessarily have to implement a test_infer_shape method or is this automated by inheriting from InferShapeTester? I see that some tests implement this method and some dont. not sure what is the best practice here.

brandonwillard · 2022-01-12T21:55:27Z

tests.unittest_tools.InferShapeTester to create automate the Op.infer_shape tests.

Could you elaborate on this part? Do I necessarily have to implement a test_infer_shape method or is this automated by inheriting from InferShapeTester? I see that some tests implement this method and some dont. not sure what is the best practice here.

Yeah, you need to create a subclass and then call the method(s) it provides. It can be a useful tool, but, if you want to make custom tests, that's also fine.

zoj613 · 2022-01-12T21:56:57Z

tests.unittest_tools.InferShapeTester to create automate the Op.infer_shape tests.

Could you elaborate on this part? Do I necessarily have to implement a test_infer_shape method or is this automated by inheriting from InferShapeTester? I see that some tests implement this method and some dont. not sure what is the best practice here.

Yeah, you need to create a subclass and then call the method(s) it provides. It can be a useful tool, but, if you want to make custom tests, that's also fine.

as in _compile_and_check?

aesara/tensor/nlinalg.py

brandonwillard · 2022-01-15T20:24:27Z

tests/tensor/test_nlinalg.py

+        self.rng = np.random.default_rng(utt.fetch_seed())
+        self.op = matmul
+        self.op_class = MatMul


This use of a shared RNG state induces test method order dependence; instead, if you create and seed the RNG objects within each independent test unit, they can be run in any order (and in parallel) and produce consistent results.

Just to be clear: by "test unit" you mean each independent test_* method?

Just to be clear: by "test unit" you mean each independent test_* method?

Yes

It looks like all the tests are sharing the same class-level RNG object. As I mentioned earlier, this will make the results order-dependent. We need to construct the RNG object within each individual test in order to avoid that.

N.B. A fixture could be used if you don't want to copy-paste the RNG construction code each time.

aesara/tensor/nlinalg.py

Sayam753 · 2022-06-21T15:28:40Z

@zoj613 is there anything left to be completed in this PR?

The PR #808 is blocked to use matmul operation (implemented in this PR).

stale

brandonwillard · 2022-06-21T19:12:16Z

It looks like this can't be rebased by maintainers, and we need to (squash and) rebase to make sure it passes with all the recent changes—especially the shape-inference-related ones. @zoj613, is "Allow edits and access to secrets by maintainers" not enabled/checked?

Sayam753 · 2022-07-07T07:45:31Z

@brandonwillard is this PR ready to be merged?

brandonwillard

Some minor changes are needed; otherwise, it looks good.

aesara/tensor/__init__.py

aesara/tensor/nlinalg.py

brandonwillard · 2022-07-11T23:55:55Z

tests/tensor/test_nlinalg.py

+        self.rng = np.random.default_rng(utt.fetch_seed())
+        self.op = matmul
+        self.op_class = MatMul


It looks like all the tests are sharing the same class-level RNG object. As I mentioned earlier, this will make the results order-dependent. We need to construct the RNG object within each individual test in order to avoid that.

N.B. A fixture could be used if you don't want to copy-paste the RNG construction code each time.

tests/tensor/test_nlinalg.py

purna135 · 2022-07-20T15:46:13Z

Thank you so much for this PR, @zoj613! I'm keeping an eye on it, and it appears that some minor changes are required before this can be merged. Are you able to push this PR to the boundary line? I have a strong dependency on this PR.

brandonwillard · 2022-07-23T20:35:26Z

I've just noticed that numpy.matmul is essentially a ufunc of dot. With that in mind, the right way to do this is to first close #695 and use that functionality to implement this.

In the meantime, I've updated the tests and docstrings.

purna135 · 2022-07-30T03:39:26Z

Hello, @brandonwillard. Is this PR now ready to be merged?
I have some PR that are waiting for this PR to be merged because my work is dependent on matrix_inverse at #808, which is once again blocked by this PR : (

brandonwillard

We can merge in order to move #808 along, but we need to keep the associated issue open until at least a gradient implementation is provided.

Preferably, we should have both a gradient and non-Python implementation for the addition of an Op, and, in this case, #757 is where our efforts need to be focused to accomplish that.

zoj613 · 2022-07-30T21:53:25Z

I am looking into adding a grad/L_op/R_op method for this. It appears that it won't be as straight forward as re-using one for the Dot Op since matmul has slight differences in behavior, especially for dimensions > 2.

brandonwillard · 2022-07-30T21:55:08Z

I am looking into adding a grad/L_op/R_op method for this. It appears that it won't be as straight forward as re-using one for the Dot Op since matmul has slight differences in behavior, especially for dimensions > 2.

Yeah, I realized that the logic for doing that is a slightly specialized form of the logic we need in #757, so it's better that we focus our efforts there.

zoj613 added the enhancement New feature or request label Jan 11, 2022

zoj613 force-pushed the matmul branch from 4699a8a to 7199d6c Compare January 11, 2022 23:23

zoj613 requested a review from brandonwillard January 11, 2022 23:24

zoj613 requested review from ricardoV94 and rlouf January 11, 2022 23:28

brandonwillard added the new operator label Jan 12, 2022

brandonwillard reviewed Jan 12, 2022

View reviewed changes

brandonwillard requested changes Jan 12, 2022

View reviewed changes

aesara/tensor/__init__.py Outdated Show resolved Hide resolved

aesara/tensor/nlinalg.py Outdated Show resolved Hide resolved

aesara/tensor/nlinalg.py Outdated Show resolved Hide resolved

brandonwillard reviewed Jan 12, 2022

View reviewed changes

tests/tensor/test_nlinalg.py Outdated Show resolved Hide resolved

zoj613 changed the title ~~WIP: Implement aesara.tensor.matmul~~ WIP: Implement aesara.tensor.matmul Jan 12, 2022

zoj613 force-pushed the matmul branch from caac71b to 78cf1be Compare January 12, 2022 21:46

zoj613 force-pushed the matmul branch from 78cf1be to 3e168d0 Compare January 12, 2022 23:45

rlouf reviewed Jan 13, 2022

View reviewed changes

aesara/tensor/nlinalg.py Outdated Show resolved Hide resolved

aesara/tensor/nlinalg.py Outdated Show resolved Hide resolved

aesara/tensor/nlinalg.py Outdated Show resolved Hide resolved

zoj613 force-pushed the matmul branch 3 times, most recently from 8ad2797 to fcb2d5c Compare January 15, 2022 09:52

brandonwillard reviewed Jan 15, 2022

View reviewed changes

aesara/tensor/nlinalg.py Outdated Show resolved Hide resolved

brandonwillard previously requested changes Jan 15, 2022

View reviewed changes

rlouf reviewed Jan 16, 2022

View reviewed changes

aesara/tensor/nlinalg.py Outdated Show resolved Hide resolved

zoj613 marked this pull request as ready for review January 16, 2022 08:16

zoj613 changed the title ~~WIP: Implement aesara.tensor.matmul~~ Implement aesara.tensor.matmul Jan 16, 2022

zoj613 force-pushed the matmul branch 2 times, most recently from 701a93e to a336429 Compare January 21, 2022 07:21

Sayam753 mentioned this pull request Feb 26, 2022

Generalize matrix_inverse to work beyond 2d arrays #808

Draft

6 tasks

zoj613 force-pushed the matmul branch from a336429 to 5abde4d Compare June 21, 2022 19:19

brandonwillard added Op implementation Involves the implementation of an Op and removed new Op labels Jul 11, 2022

brandonwillard requested changes Jul 11, 2022

View reviewed changes

brandonwillard force-pushed the matmul branch from 5abde4d to 7a058dc Compare July 23, 2022 21:00

brandonwillard added NumPy compatibility tensor algebra Relates to our use and representations of tensor algebra labels Jul 23, 2022

brandonwillard force-pushed the matmul branch 2 times, most recently from b9e3b2c to 7e0f94c Compare July 23, 2022 21:03

Add an Op for numpy.matmul

b60030f

brandonwillard force-pushed the matmul branch from 7e0f94c to b60030f Compare July 24, 2022 19:51

brandonwillard approved these changes Jul 30, 2022

View reviewed changes

brandonwillard merged commit 2246504 into aesara-devs:main Jul 30, 2022

brandonwillard mentioned this pull request Jul 30, 2022

Implement aesara.tensor.matmul #488

Open

zoj613 deleted the matmul branch July 30, 2022 21:58

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement `aesara.tensor.matmul` #744

Implement `aesara.tensor.matmul` #744

zoj613 commented Jan 11, 2022 •

edited

Loading

codecov bot commented Jan 11, 2022 •

edited

Loading

brandonwillard left a comment

brandonwillard left a comment

zoj613 commented Jan 12, 2022

zoj613 commented Jan 12, 2022

brandonwillard commented Jan 12, 2022

zoj613 commented Jan 12, 2022

brandonwillard Jan 15, 2022

zoj613 Jan 15, 2022

brandonwillard Jan 15, 2022

brandonwillard Jul 11, 2022

Sayam753 commented Jun 21, 2022

brandonwillard commented Jun 21, 2022

Sayam753 commented Jul 7, 2022

brandonwillard left a comment

brandonwillard Jul 11, 2022

purna135 commented Jul 20, 2022

brandonwillard commented Jul 23, 2022 •

edited

Loading

purna135 commented Jul 30, 2022

brandonwillard left a comment

zoj613 commented Jul 30, 2022

brandonwillard commented Jul 30, 2022

Implement aesara.tensor.matmul #744

Implement aesara.tensor.matmul #744

Conversation

zoj613 commented Jan 11, 2022 • edited Loading

References

codecov bot commented Jan 11, 2022 • edited Loading

Codecov Report

brandonwillard left a comment

Choose a reason for hiding this comment

brandonwillard left a comment

Choose a reason for hiding this comment

zoj613 commented Jan 12, 2022

zoj613 commented Jan 12, 2022

brandonwillard commented Jan 12, 2022

zoj613 commented Jan 12, 2022

brandonwillard Jan 15, 2022

Choose a reason for hiding this comment

zoj613 Jan 15, 2022

Choose a reason for hiding this comment

brandonwillard Jan 15, 2022

Choose a reason for hiding this comment

brandonwillard Jul 11, 2022

Choose a reason for hiding this comment

Sayam753 commented Jun 21, 2022

brandonwillard commented Jun 21, 2022

Sayam753 commented Jul 7, 2022

brandonwillard left a comment

Choose a reason for hiding this comment

brandonwillard Jul 11, 2022

Choose a reason for hiding this comment

purna135 commented Jul 20, 2022

brandonwillard commented Jul 23, 2022 • edited Loading

purna135 commented Jul 30, 2022

brandonwillard left a comment

Choose a reason for hiding this comment

zoj613 commented Jul 30, 2022

brandonwillard commented Jul 30, 2022

Implement `aesara.tensor.matmul` #744

Implement `aesara.tensor.matmul` #744

zoj613 commented Jan 11, 2022 •

edited

Loading

codecov bot commented Jan 11, 2022 •

edited

Loading

brandonwillard commented Jul 23, 2022 •

edited

Loading