Implement Dot and BatchedDot in PyTensor #878

HangenYuu · 2024-07-03T07:42:57Z

Description

Implemented the PyTorch link and unit tests for the math Op Dot and blas Op BatchedDot. Did not touch on BatchedDot or Dot for sparse matrices.

Progress

Dot
BatchedDot

Related Issue

Closes #
Related to Implement all Ops in PyTorch (help welcome!) #821

Checklist

Checked that the pre-commit linting/style checks pass
Included tests that prove the fix is effective or that the new feature works
Added necessary documentation (docstrings and/or example notebooks)
If you are a pro: each commit corresponds to a relevant logical change

Type of change

pytensor/link/pytorch/dispatch/nlinalg.py

ricardoV94

Well that was a round about trip :D Thanks!

ricardoV94 · 2024-07-08T12:39:29Z

pytensor/link/pytorch/dispatch/nlinalg.py

+
+
+@pytorch_funcify.register(Dot)
+def pytorch_funcify_Dot(op, **kwargs):


You have to import this file from pytorch.dispatch.__init__ for it to be registered (the test is failing in the CI). But Dot is not defined in nlinalg, so we should put it in dispatch/match.py? Same for the test

I based it off the JAX link. If you take a look at pytensor/link/jax/dispatch/nlinalg.py you will see Max, Argmax, and Dot Ops from math in there. Do you want me to separate them out for JAX too?

I can also put the Argmax I am implementing in pytorch/dispatch/math.py.

Yeah in general we want to keep it more or less mirrored with the file structure where they are defined. Although our tensor/basic.py and tensor/math.py are in need of being split of as they have way too many lines

ricardoV94 · 2024-07-08T15:17:28Z

BatchedDot should be pretty simple as well, it's a matmul of 3d tensors without allowing broadcasting on the leading dimension, so with a check that a.shape[0] == b.shape[0]. Wanna give it a go? You should be able to trigger it if you test tensor(shape=(5, 3, 2)) @ tensor(shape=(5, 2, 4)) (pseudo-code)

HangenYuu · 2024-07-10T03:33:54Z

Yep, torch.argmax also does not allow multiple partial axes.

HangenYuu · 2024-07-10T08:15:30Z

@ricardoV94 BatchedDot is done. I will do Max and Argmax next. They are tougher nuts to crack.

ricardoV94 · 2024-07-10T08:23:31Z

Isn't Max done already? Should be like Sum/All, ... which we did already. Argmax just needs some ravelling and tranposing. You should be able to copy the logic inside perform

ricardoV94 · 2024-07-10T08:24:03Z

pytensor/link/jax/dispatch/blas.py

+def jax_funcify_BatchedDot(op, **kwargs):
+    def batched_dot(a, b):
+        if a.shape[0] != b.shape[0]:
+            raise TypeError("Shapes must match in the 0-th dimension")


Suggested change

raise TypeError("Shapes must match in the 0-th dimension")

raise TypeError("Shapes must match along the first dimension of BatchedDot")

ricardoV94 · 2024-07-10T08:26:08Z

Can you split the JAX changes into a separate PR? It's better have PRs atomic as it makes it easier to review. Sometimes it's okay to have multiple functionality in a PR but then you have to respect this part of the checklist:

If you are a pro: each commit corresponds to a relevant logical change

ricardoV94 · 2024-07-10T08:27:54Z

pytensor/link/jax/dispatch/math.py

+
+
+@jax_funcify.register(Max)
+def jax_funcify_Max(op, **kwargs):


We should have the dispatch for the other CAReduce, like All, Any... on the same file as Max

HangenYuu · 2024-07-10T08:59:45Z

Can you split the JAX changes into a separate PR? It's better have PRs atomic as it makes it easier to review. Sometimes it's okay to have multiple functionality in a PR but then you have to respect this part of the checklist:

Okay I moved it to #913

HangenYuu · 2024-07-15T02:21:43Z

@ricardoV94 For this one, I will stop at BatchedDot.

ricardoV94 · 2024-07-15T11:31:25Z

tests/link/pytorch/test_blas.py

+    opts = RewriteDatabaseQuery(include=[None], exclude=["cxx_only", "BlasOpt"])
+    pytorch_mode = Mode(PytorchLinker(), opts)
+    pytensor_pytorch_fn = function(fgraph.inputs, fgraph.outputs, mode=pytorch_mode)


This does the same?

Suggested change

opts = RewriteDatabaseQuery(include=[None], exclude=["cxx_only", "BlasOpt"])

pytorch_mode = Mode(PytorchLinker(), opts)

pytensor_pytorch_fn = function(fgraph.inputs, fgraph.outputs, mode=pytorch_mode)

pytorch_mode_no_rewrites = Mode(PytorchLinker(), None)

pytensor_pytorch_fn = function(fgraph.inputs, fgraph.outputs, mode= pytorch_mode_no_rewrites)

But if I am not mistaken compare_pytorch_and_py returns the torch function, so you could just reuse it?

codecov · 2024-07-17T02:51:47Z

Codecov Report

Attention: Patch coverage is 85.71429% with 3 lines in your changes missing coverage. Please review.

Project coverage is 81.40%. Comparing base (426931b) to head (f459866).
Report is 94 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/link/pytorch/dispatch/blas.py	80.00%	2 Missing ⚠️
pytensor/link/pytorch/dispatch/math.py	87.50%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #878   +/-   ##
=======================================
  Coverage   81.40%   81.40%           
=======================================
  Files         173      175    +2     
  Lines       46914    46934   +20     
  Branches    11426    11427    +1     
=======================================
+ Hits        38188    38205   +17     
- Misses       6544     6547    +3     
  Partials     2182     2182

Files with missing lines	Coverage Δ
pytensor/link/pytorch/dispatch/__init__.py	`100.00% <100.00%> (ø)`
pytensor/link/pytorch/dispatch/math.py	`87.50% <87.50%> (ø)`
pytensor/link/pytorch/dispatch/blas.py	`80.00% <80.00%> (ø)`

ricardoV94 · 2024-07-17T06:22:57Z

tests/link/pytorch/test_blas.py

+    pytorch_mode_no_rewrites = Mode(PytorchLinker(), None)
+    pytensor_pytorch_fn.mode = pytorch_mode_no_rewrites


This is not a thing you can do (or rather has no effect). Once a function it's compiled that's it, the mode plays no role anymore

Suggested change

pytorch_mode_no_rewrites = Mode(PytorchLinker(), None)

pytensor_pytorch_fn.mode = pytorch_mode_no_rewrites

ricardoV94 · 2024-07-17T06:23:46Z

tests/link/pytorch/test_blas.py

+    a.tag.test_value = (
+        np.linspace(-1, 1, 10 * 5 * 3).astype(config.floatX).reshape((10, 5, 3))
+    )


We are getting rid of the test_value machinery. Just pass these directly to the test function, no point in putting them in the tag to then retrieve it again

ricardoV94

Looks good, just a nit if you want to address

ricardoV94 · 2024-07-17T10:32:41Z

tests/link/pytorch/test_blas.py

+def test_pytorch_BatchedDot():
+    # tensor3 . tensor3
+    a = tensor3("a")
+    A = np.linspace(-1, 1, 10 * 5 * 3).astype(config.floatX).reshape((10, 5, 3))


Nit: A more conventional name would be a_test, b_test for variables a and b

ricardoV94 · 2024-07-18T07:07:15Z

GitHub says the branch has conflicts with main. Can you update?

HangenYuu · 2024-07-18T08:59:45Z

GitHub says the branch has conflicts with main. Can you update?

Done.

Added PyTorch link and unit tests for normal dot

ffad937

ricardoV94 reviewed Jul 3, 2024

View reviewed changes

pytensor/link/pytorch/dispatch/nlinalg.py Outdated Show resolved Hide resolved

ricardoV94 mentioned this pull request Jul 4, 2024

Implement all Ops in PyTorch (help welcome!) #821

Open

48 tasks

ricardoV94 added enhancement New feature or request torch PyTorch backend labels Jul 4, 2024

HangenYuu added 3 commits July 6, 2024 17:30

Changed implementation of dot. Renamed tests

5121a85

Changed dot implementation

2721c5a

Reverted logic to correct scope for math.dot

03bb3a8

ricardoV94 approved these changes Jul 8, 2024

View reviewed changes

ricardoV94 reviewed Jul 8, 2024

View reviewed changes

HangenYuu changed the title ~~Added PyTorch link and unit tests for normal dot~~ Added PyTorch link and unit tests for blas & math Ops Jul 10, 2024

HangenYuu mentioned this pull request Jul 10, 2024

Added PyTorch link and unit tests for argmax #900

Closed

11 tasks

ricardoV94 reviewed Jul 10, 2024

View reviewed changes

HangenYuu mentioned this pull request Jul 10, 2024

Reorganized JAX link folder structure #913

Merged

11 tasks

HangenYuu changed the title ~~Added PyTorch link and unit tests for blas & math Ops~~ Added PyTorch link and unit tests for math.Dot and blas.BatchedDot Jul 15, 2024

Reverted folder structure and added BatchedDot

2cf0ed2

HangenYuu force-pushed the torch_dot branch from c643494 to 2cf0ed2 Compare July 15, 2024 02:28

HangenYuu and others added 4 commits July 15, 2024 09:28

Merge branch 'main' into torch_dot

fcb3b79

Fixed minor typo in test naming

307a3fb

Merge branch 'torch_dot' of github.com:HangenYuu/pytensor into torch_dot

e2500bf

Fixed __init__.py file for tests to run

143a75a

ricardoV94 reviewed Jul 15, 2024

View reviewed changes

ricardoV94 changed the title ~~Added PyTorch link and unit tests for math.Dot and blas.BatchedDot~~ Implement Dot and BatchedDot in PyTensor Jul 15, 2024

Rewrite test to reuse pytorch function

2d74b31

ricardoV94 requested changes Jul 17, 2024

View reviewed changes

Removed get_test_value

4deea70

ricardoV94 approved these changes Jul 17, 2024

View reviewed changes

Changed variable names

cab9db8

Merge branch 'main' into torch_dot

f459866

ricardoV94 merged commit 6ad1c5c into pymc-devs:main Jul 18, 2024
58 of 59 checks passed

HangenYuu deleted the torch_dot branch July 22, 2024 02:15

Ch0ronomato pushed a commit to Ch0ronomato/pytensor that referenced this pull request Aug 15, 2024

Implement Dot and BatchedDot in PyTensor (pymc-devs#878)

79232b2

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement Dot and BatchedDot in PyTensor #878

Implement Dot and BatchedDot in PyTensor #878

HangenYuu commented Jul 3, 2024 •

edited

Loading

ricardoV94 left a comment

ricardoV94 Jul 8, 2024 •

edited

Loading

HangenYuu Jul 8, 2024

HangenYuu Jul 8, 2024

ricardoV94 Jul 8, 2024

ricardoV94 commented Jul 8, 2024 •

edited

Loading

HangenYuu commented Jul 10, 2024

HangenYuu commented Jul 10, 2024

ricardoV94 commented Jul 10, 2024

ricardoV94 Jul 10, 2024

ricardoV94 commented Jul 10, 2024 •

edited

Loading

ricardoV94 Jul 10, 2024

HangenYuu commented Jul 10, 2024 •

edited

Loading

HangenYuu commented Jul 15, 2024

ricardoV94 Jul 15, 2024

ricardoV94 Jul 15, 2024

codecov bot commented Jul 17, 2024 •

edited

Loading

ricardoV94 Jul 17, 2024

ricardoV94 Jul 17, 2024

ricardoV94 left a comment

ricardoV94 Jul 17, 2024

ricardoV94 commented Jul 18, 2024

HangenYuu commented Jul 18, 2024



		@pytorch_funcify.register(Dot)
		def pytorch_funcify_Dot(op, **kwargs):

	raise TypeError("Shapes must match in the 0-th dimension")
	raise TypeError("Shapes must match along the first dimension of BatchedDot")



		@jax_funcify.register(Max)
		def jax_funcify_Max(op, **kwargs):

		pytorch_mode_no_rewrites = Mode(PytorchLinker(), None)
		pytensor_pytorch_fn.mode = pytorch_mode_no_rewrites

Implement Dot and BatchedDot in PyTensor #878

Implement Dot and BatchedDot in PyTensor #878

Conversation

HangenYuu commented Jul 3, 2024 • edited Loading

Description

Related Issue

Checklist

Type of change

ricardoV94 left a comment

Choose a reason for hiding this comment

ricardoV94 Jul 8, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ricardoV94 commented Jul 8, 2024 • edited Loading

HangenYuu commented Jul 10, 2024

HangenYuu commented Jul 10, 2024

ricardoV94 commented Jul 10, 2024

Choose a reason for hiding this comment

ricardoV94 commented Jul 10, 2024 • edited Loading

Choose a reason for hiding this comment

HangenYuu commented Jul 10, 2024 • edited Loading

HangenYuu commented Jul 15, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov bot commented Jul 17, 2024 • edited Loading

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ricardoV94 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ricardoV94 commented Jul 18, 2024

HangenYuu commented Jul 18, 2024

HangenYuu commented Jul 3, 2024 •

edited

Loading

ricardoV94 Jul 8, 2024 •

edited

Loading

ricardoV94 commented Jul 8, 2024 •

edited

Loading

ricardoV94 commented Jul 10, 2024 •

edited

Loading

HangenYuu commented Jul 10, 2024 •

edited

Loading

codecov bot commented Jul 17, 2024 •

edited

Loading