Add docs on implementing Pytorch Ops (and CumOp) #837

HarshvirSandhu · 2024-06-20T16:57:27Z

Description

This PR can be used as an example for implementing Ops in PyTorch

Related Issue

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)

Type of change

cc @ricardoV94

ricardoV94 · 2024-06-20T17:16:39Z

pytensor/link/pytorch/dispatch/extra_ops.py

+    dim = op.axis
+    mode = op.mode
+
+    def cumop(x, dim=dim, mode=mode):


This is not needed, the returned functions are never called by the user

Suggested change

def cumop(x, dim=dim, mode=mode):

def cumop(x):

ricardoV94 · 2024-06-20T17:17:28Z

tests/link/pytorch/test_extra_ops.py

+    # Create test value tag for a
+    a.tag.test_value = np.arange(9, dtype=config.floatX).reshape((3, 3))


No need for test values and tags. We're planning to deprecate that functionality as well

ricardoV94 · 2024-06-20T17:17:55Z

tests/link/pytorch/test_extra_ops.py

+    # For the second mode of CumOp
+    out = pt.cumprod(a, axis=1)
+    fgraph = FunctionGraph([a], [out])
+    compare_pytorch_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])


Here just pass the test values (instead of adding them as tags and then retrieving them)

ricardoV94 · 2024-06-20T17:19:20Z

tests/link/pytorch/test_extra_ops.py

+    a.tag.test_value = np.arange(9, dtype=config.floatX).reshape((3, 3))
+
+    # Create the output variable
+    out = pt.cumsum(a, axis=0)


Test axis=None and axis=tuple(...) if supported by the original Op. If tuple is allowed make sure you have more dimensions (say 3) and only ask for a subset (say 2) of them in the axis. This is to make sure you test something that is different than axis=None or axis=int.

The axis can be parametrized (prod and add as well) instead of adding more conditions inside the test

Tried this on the original Op. axis=tuple(...) does not work and gives a TypeError
axis=None gives the output as a 1-D array

The Op __init__ doesn't seem to check explicitly for axes but it does assume it is either None or an int. Can we add a check and raise an explicit ValueError if it's not either?

Checked again, there is no error if we use axis=(0), pytorch also returns the same output.
The error only comes when there are more than 1 elements in the tuple (Even np.cumsum gives TypeError in this case).

We could try adding a check and raise, but would that be needed in other Op implementations?
Since this would be used as an example, it might be complicated if a check and raise is not needed for other implementations.

(0) is 0, not a tuple with a 0 inside it, it would have to be (0,) to be a tuple with a single element inside. Does it work with (0,)?

No, it gives a TypeError

Which is fine but probably gives a typeerror in an obscure place. We should raise already in the init method of the Op to save people time

ricardoV94 · 2024-06-20T17:21:56Z

Can you extend the example in the documentation page on implementing custom JAX/NUMBA Ops to mention PyTorch and include this example as well?

Perhaps you can use some fancy tab to select among the different modes in the same documentation page. Is that supported @OriolAbril ?

OriolAbril · 2024-06-20T18:45:35Z

Not here as of now, you'd have to add an extra extension for tabs. If you'll only want tabs, then it is probably best to use https://sphinx-tabs.readthedocs.io/en/latest/, if using things like grids, dropdowns, icons... somewhere else in addition to tabs here seems a future possibility then https://sphinx-design.readthedocs.io/en/sbt-theme/ is probably best. Both should only require being added as dependencies to the doc env and adding them to the extensions varialbe in conf.py, no further configuration

ricardoV94 · 2024-06-20T21:16:30Z

Thanks @OriolAbril either of those seems perfect. Any preference?

codecov · 2024-06-22T11:07:03Z

Codecov Report

Attention: Patch coverage is 63.15789% with 7 lines in your changes missing coverage. Please review.

Project coverage is 80.97%. Comparing base (320bac4) to head (a6e6bd8).
Report is 22 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #837      +/-   ##
==========================================
+ Coverage   80.87%   80.97%   +0.10%     
==========================================
  Files         168      170       +2     
  Lines       46950    47044      +94     
  Branches    11472    11504      +32     
==========================================
+ Hits        37972    38096     +124     
+ Misses       6766     6734      -32     
- Partials     2212     2214       +2

Files	Coverage Δ
pytensor/link/pytorch/dispatch/__init__.py	`100.00% <100.00%> (ø)`
pytensor/tensor/extra_ops.py	`88.63% <100.00%> (+0.03%)`	⬆️
pytensor/link/pytorch/dispatch/extra_ops.py	`56.25% <56.25%> (ø)`

... and 20 files with indirect coverage changes

ricardoV94 · 2024-06-22T15:32:05Z

pytensor/link/pytorch/dispatch/extra_ops.py

+    dim = op.axis
+    mode = op.mode
+
+    def cumop(x, dim=dim):


Looks good, just no need for any kwargs. The function will only ever receive the node inputs

Suggested change

def cumop(x, dim=dim):

def cumop(x):

OriolAbril · 2024-06-22T16:31:34Z

Thanks @OriolAbril either of those seems perfect. Any preference?

I use sphinx-design more because I use its other features

ricardoV94 · 2024-06-22T15:33:31Z

tests/link/pytorch/test_extra_ops.py

+    # Create a symbolic input for the first input of `CumOp`
+    a = pt.matrix("a")
+
+    # Create test value tag for a


Suggested change

# Create test value tag for a

# Create test value

ricardoV94 · 2024-06-23T09:53:21Z

tests/link/pytorch/test_extra_ops.py

+    a.tag.test_value = np.arange(9, dtype=config.floatX).reshape((3, 3))
+
+    # Create the output variable
+    out = pt.cumsum(a, axis=0)


Which is fine but probably gives a typeerror in an obscure place. We should raise already in the init method of the Op to save people time

ricardoV94 · 2024-06-23T21:00:33Z

pytensor/tensor/extra_ops.py

@@ -283,8 +283,11 @@ class CumOp(COp):
    def __init__(self, axis: int | None = None, mode="add"):
        if mode not in ("add", "mul"):
            raise ValueError(f'{type(self).__name__}: Unknown mode "{mode}"')
-        self.axis = axis
-        self.mode = mode
+        if isinstance(axis, int) or axis is None:


Nitpick, it's more common to just check and raise than indenting the "correct code" and raising otherwise

Suggested change

if isinstance(axis, int) or axis is None:

if not (isinstance(axis, int) or axis is None):

# raise error

# usual code

That's how the error check above for the mode is structured as well

ricardoV94 · 2024-06-23T22:11:33Z

doc/extending/creating_a_numba_jax_op.rst

-       return res if n_outs > 1 else res[0]
+.. tab-set::
+
+        .. tab-item:: JAX/Numba     


This is not correct for Numba, can you leave it as a separate tab with [in progress] text (and open an issue) or check the source code of the Numba implementation if you want to do it correctly?

This probably applies to all the tabbed sections, no reason to combine jax and numba, and the pre-existing snippets were JAX specific

I was adding a separate tab for numba and found this comment. Is there anything that should be changed in numba_funcify_DimShuffle?

pytensor/pytensor/link/numba/dispatch/elemwise.py

Lines 688 to 690 in 7159215

# FIXME: Numba's `array.reshape` only accepts C arrays.

res_reshape = np.reshape(np.ascontiguousarray(x), new_shape)

Not in the context of this PR, but we should open an issue here to check if that's still a problem in the newer versions of numba. Could you do that?

OriolAbril

all tabs look rendered correctly, only left a comment so cross references to other libraries actually work

OriolAbril · 2024-06-24T17:17:19Z

doc/extending/creating_a_numba_jax_op.rst

 function that performs exactly the same computations as the :class:`Op`. For
 example, the :class:`Eye` operator has a JAX equivalent: :func:`jax.numpy.eye`
-(see `the documentation <https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.eye.html?highlight=eye>`_).
+(see `the documentation <https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.eye.html?highlight=eye>`_) and a Pytorch equivalent :func:`torch.eye` (see `documentation <https://pytorch.org/docs/stable/generated/torch.eye.html>`_).


This looks like this:

which is quite the weird pattern for docs, especially given jax.numpy.eye and torch.eye are already using the correct cross-referencing syntax. I would remove the manual links and use the cross-references. That is, leaving only this:

Suggested change

(see `the documentation <https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.eye.html?highlight=eye>`_) and a Pytorch equivalent :func:`torch.eye` (see `documentation <https://pytorch.org/docs/stable/generated/torch.eye.html>`_).

and a Pytorch equivalent :func:`torch.eye`.

And doing two more changes to conf.py. First add sphinx.ext.intersphinx to the list of extensions. It is part of the main sphinx library so no need to add any extra dependency to the env file. Add

intersphinx_mapping = { "jax": ("https://jax.readthedocs.io/en/latest", None), "numpy": ("https://numpy.org/doc/stable", None), "torch": ("https://pytorch.org/docs/stable", None), }

with that, the jax.numpy.eye and torch.eye will still be formatted as monospaced text but no longer be pink, they'll be blue and be clickable links to their respective API pages.

Add Pytorch support for Cum Op

f007c0d

ricardoV94 reviewed Jun 20, 2024

View reviewed changes

Modify test for Cum op

6ae355f

ricardoV94 reviewed Jun 22, 2024

View reviewed changes

ricardoV94 reviewed Jun 23, 2024

View reviewed changes

Raise TypeError if axis not int or None

33463fc

ricardoV94 reviewed Jun 23, 2024

View reviewed changes

HarshvirSandhu added 2 commits June 24, 2024 03:26

Fix init method of CumOp

bf905cb

Extend tutorial on documentation for Pytorch

debc3e0

ricardoV94 reviewed Jun 23, 2024

View reviewed changes

Add tab for Numba

2bc7ddc

ricardoV94 approved these changes Jun 24, 2024

View reviewed changes

ricardoV94 requested a review from OriolAbril June 24, 2024 17:05

OriolAbril reviewed Jun 24, 2024

View reviewed changes

Add intersphinx mapping

ec87e4e

ricardoV94 added enhancement New feature or request torch PyTorch backend labels Jun 28, 2024

ricardoV94 mentioned this pull request Jun 28, 2024

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

Open

48 tasks

Parametrize dtype

a6e6bd8

ricardoV94 approved these changes Jun 28, 2024

View reviewed changes

ricardoV94 requested a review from OriolAbril June 28, 2024 11:34

ricardoV94 merged commit 781073b into pymc-devs:main Jul 4, 2024
56 of 57 checks passed

ricardoV94 changed the title ~~Add Pytorch support for Cum Op~~ Add docs on implementing Pytorch Ops (and CumOp) Jul 4, 2024

ricardoV94 added the docs label Jul 4, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add docs on implementing Pytorch Ops (and CumOp) #837

Add docs on implementing Pytorch Ops (and CumOp) #837

HarshvirSandhu commented Jun 20, 2024

ricardoV94 Jun 20, 2024

ricardoV94 Jun 20, 2024

ricardoV94 Jun 20, 2024

ricardoV94 Jun 20, 2024 •

edited

Loading

HarshvirSandhu Jun 22, 2024

ricardoV94 Jun 22, 2024

HarshvirSandhu Jun 22, 2024

OriolAbril Jun 22, 2024

HarshvirSandhu Jun 22, 2024

ricardoV94 Jun 23, 2024

ricardoV94 commented Jun 20, 2024 •

edited

Loading

OriolAbril commented Jun 20, 2024

ricardoV94 commented Jun 20, 2024

codecov bot commented Jun 22, 2024 •

edited

Loading

ricardoV94 Jun 22, 2024

OriolAbril commented Jun 22, 2024

ricardoV94 Jun 22, 2024

ricardoV94 Jun 23, 2024

ricardoV94 Jun 23, 2024 •

edited

Loading

ricardoV94 Jun 23, 2024

HarshvirSandhu Jun 24, 2024

ricardoV94 Jun 24, 2024 •

edited

Loading

OriolAbril left a comment

OriolAbril Jun 24, 2024

		# Create test value tag for a
		a.tag.test_value = np.arange(9, dtype=config.floatX).reshape((3, 3))

-        if isinstance(axis, int) or axis is None:
+        if not (isinstance(axis, int) or axis is None):
+            # raise error
+        # usual code


	# FIXME: Numba's `array.reshape` only accepts C arrays.
	res_reshape = np.reshape(np.ascontiguousarray(x), new_shape)

	(see `the documentation <https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.eye.html?highlight=eye>`_) and a Pytorch equivalent :func:`torch.eye` (see `documentation <https://pytorch.org/docs/stable/generated/torch.eye.html>`_).
	and a Pytorch equivalent :func:`torch.eye`.

Add docs on implementing Pytorch Ops (and CumOp) #837

Add docs on implementing Pytorch Ops (and CumOp) #837

Conversation

HarshvirSandhu commented Jun 20, 2024

Description

Related Issue

Checklist

Type of change

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ricardoV94 Jun 20, 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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ricardoV94 commented Jun 20, 2024 • edited Loading

OriolAbril commented Jun 20, 2024

ricardoV94 commented Jun 20, 2024

codecov bot commented Jun 22, 2024 • edited Loading

Codecov Report

Choose a reason for hiding this comment

OriolAbril commented Jun 22, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ricardoV94 Jun 23, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ricardoV94 Jun 24, 2024 • edited Loading

Choose a reason for hiding this comment

OriolAbril left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ricardoV94 Jun 20, 2024 •

edited

Loading

ricardoV94 commented Jun 20, 2024 •

edited

Loading

codecov bot commented Jun 22, 2024 •

edited

Loading

ricardoV94 Jun 23, 2024 •

edited

Loading

ricardoV94 Jun 24, 2024 •

edited

Loading