Implements shape Ops and MakeVector in PyTorch #926

twaclaw · 2024-07-12T21:52:19Z

Description

Implements

Shape
Shape_i
Reshape
SpecifyShape
Unbroadcast
MakeVector

Related Issue

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

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

codecov · 2024-07-12T22:16:12Z

Codecov Report

Attention: Patch coverage is 80.85106% with 9 lines in your changes missing coverage. Please review.

Project coverage is 81.40%. Comparing base (72c6a81) to head (0786f2c).
Report is 98 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/link/pytorch/dispatch/shape.py	74.28%	8 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #926      +/-   ##
==========================================
+ Coverage   81.38%   81.40%   +0.01%     
==========================================
  Files         172      173       +1     
  Lines       46868    46914      +46     
  Branches    11423    11426       +3     
==========================================
+ Hits        38145    38188      +43     
- Misses       6540     6544       +4     
+ Partials     2183     2182       -1

Files with missing lines	Coverage Δ
pytensor/link/pytorch/dispatch/__init__.py	`100.00% <100.00%> (ø)`
pytensor/link/pytorch/dispatch/basic.py	`89.74% <100.00%> (+4.44%)`	⬆️
pytensor/link/pytorch/dispatch/shape.py	`74.28% <74.28%> (ø)`

... and 4 files with indirect coverage changes

ricardoV94 · 2024-07-13T15:04:55Z

pytensor/link/pytorch/dispatch/shape.py

+@pytorch_funcify.register(Reshape)
+def pytorch_funcify_Reshape(op, node, **kwargs):
+    shape = node.inputs[1]
+
+    def reshape(x, shape=shape):
+        return torch.reshape(x, tuple(shape))
+
+    return reshape


We have to use the runtime shape, since it is not always a constant.

Suggested change

@pytorch_funcify.register(Reshape)

def pytorch_funcify_Reshape(op, node, **kwargs):

shape = node.inputs[1]

def reshape(x, shape=shape):

return torch.reshape(x, tuple(shape))

return reshape

@pytorch_funcify.register(Reshape)

def pytorch_funcify_Reshape(op, node, **kwargs):

def reshape(x, shape):

return torch.reshape(x, tuple(shape))

return reshape

ricardoV94 · 2024-07-13T15:05:58Z

tests/link/pytorch/test_shape.py

+    compare_pytorch_and_py(x_fg, [np.r_[1.0, 2.0, 3.0, 4.0].astype(config.floatX)])
+
+
+def test_pytorch_Reshape_shape_graph_input():


Suggested change

def test_pytorch_Reshape_shape_graph_input():

def test_pytorch_Reshape_dynamic():

ricardoV94 · 2024-07-13T15:07:48Z

tests/link/pytorch/test_shape.py

+    x = DeepCopyOp()(pt.as_tensor_variable(1.1))
+    x_fg = FunctionGraph([], [x])


We have to make sure DeepCopy does the expected thing. For instance here is how we know it is currently not doing the right thing in the Numba backend: #50

Is the DeepCopyOp actually in the scope of these changes?
Not sure why DeepCopy and View were being tested together with Unbroadcast (I adapted the tests from the JAX backend implementation).

You can leave then out if you didn't mean to implement it

Fwiw, i tested the torch clone and it seems like it's fine

import pytensor from pytensor import tensor as pt x = pytensor.shared(0, name="x") f = pytensor.function([], x, mode=None) f().itemset(2) assert x.get_value() == 0 f = pytensor.function([], x, mode="PYTORCH") f().apply_(lambda _: 2) assert x.get_value() == 0

ricardoV94 · 2024-07-13T15:09:44Z

pytensor/link/pytorch/dispatch/basic.py

 @singledispatch
 def pytorch_typify(data, dtype=None, **kwargs):
    r"""Convert instances of PyTensor `Type`\s to PyTorch types."""
-    return torch.as_tensor(data, dtype=dtype)
+    if data is not None:
+        return torch.as_tensor(data, dtype=dtype)
+    return None


We should dispatch on NoneType:

@pytorch_typify.register(NoneType): def pytorch_typify_None(data, **kwargs): return None

Where should this go?

In addition to condition in def pytorch_typify(data, dtype=None, **kwargs): ... ?

It should be fine in dispatch.basic.py.

If you dispatch you don't need the if, because the dispatch mechanism already chooses which function to call based on the type of data

I tried something like:

from pytensor.tensor.type_other import NoneTypeT @pytorch_typify.register(NoneTypeT) def pytorch_typify_None(data, **kwargs): return None

but the condition is still required. I checked the JAX backend for reference and there is also a similar condition there for pytorch_typify(data ...

A different issue surfaced. I don't know whether it is related to these changes or not.

Calling repeat with axis=None results in an error in ElementWise at this point; namely:

torch._dynamo.exc.InternalTorchDynamoError: 'int' object has no attribute 'shape'

the values of inputs is (tensor(3), 3, 2). The corresponding graph is shown below.

Reshape{1} [id A] <Vector(float64, shape=(?,))> 8 ├─ Alloc [id B] <Matrix(float64, shape=(?, 3))> 7 │ ├─ ExpandDims{axis=1} [id C] <Matrix(float64, shape=(?, 1))> 6 │ │ └─ Reshape{1} [id D] <Vector(float64, shape=(?,))> 5 │ │ ├─ a [id E] <Matrix(float64, shape=(?, ?))> │ │ └─ [-1] [id F] <Vector(int64, shape=(1,))> │ ├─ Mul [id G] <Scalar(int64, shape=())> 4 │ │ ├─ Shape_i{0} [id H] <Scalar(int64, shape=())> 1 │ │ │ └─ a [id E] <Matrix(float64, shape=(?, ?))> │ │ └─ Shape_i{1} [id I] <Scalar(int64, shape=())> 0 │ │ └─ a [id E] <Matrix(float64, shape=(?, ?))> │ └─ 3 [id J] <Scalar(int64, shape=())> └─ MakeVector{dtype='int64'} [id K] <Vector(int64, shape=(1,))> 3 └─ Mul [id L] <Scalar(int64, shape=())> 2 ├─ 3 [id J] <Scalar(int64, shape=())> ├─ Shape_i{0} [id H] <Scalar(int64, shape=())> 1 │ └─ ··· └─ Shape_i{1} [id I] <Scalar(int64, shape=())> 0 └─ ···

Typify should be registered on NoneType (python) not NoneTypeT (PyTensor)

Regarding your other error, there shouldn't be integers 3 and 2 as inputs, it should be tensor(2) and tensor(3).

We need to track where those come from and fix it

Probably Shape_i is not correctly implemented and is returning integers?

- Shape - Shape_i - Reshape - SpecifyShape - Unbroadcast - MakeVector

- Fixed Shape_i - Typified Python NoneType

ricardoV94 · 2024-07-17T19:05:45Z

Awesome stuff @twaclaw

ricardoV94 reviewed Jul 13, 2024

View reviewed changes

ricardoV94 mentioned this pull request Jul 15, 2024

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

Open

48 tasks

twaclaw added 2 commits July 15, 2024 19:45

Implements shape and MakeVector Ops in PyTorch

4d4abd7

- Shape - Shape_i - Reshape - SpecifyShape - Unbroadcast - MakeVector

Reworked tests in implementation of Shape in PyTorch

0e455fd

twaclaw force-pushed the implement_shape_ops_and_makevector_in_pytorch branch from bf50423 to 0e455fd Compare July 16, 2024 18:21

Fixed implementation of Shape Op in PyTorch

0786f2c

- Fixed Shape_i - Typified Python NoneType

ricardoV94 approved these changes Jul 17, 2024

View reviewed changes

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

ricardoV94 merged commit 426931b into pymc-devs:main Jul 17, 2024
58 of 59 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implements shape Ops and MakeVector in PyTorch #926

Implements shape Ops and MakeVector in PyTorch #926

twaclaw commented Jul 12, 2024

codecov bot commented Jul 12, 2024 •

edited

Loading

ricardoV94 Jul 13, 2024

ricardoV94 Jul 13, 2024

ricardoV94 Jul 13, 2024

twaclaw Jul 14, 2024

ricardoV94 Jul 14, 2024

Ch0ronomato Jul 16, 2024 •

edited

Loading

ricardoV94 Jul 13, 2024

twaclaw Jul 14, 2024

ricardoV94 Jul 15, 2024

twaclaw Jul 16, 2024

twaclaw Jul 16, 2024 •

edited

Loading

ricardoV94 Jul 17, 2024

ricardoV94 Jul 17, 2024 •

edited

Loading

ricardoV94 Jul 17, 2024

ricardoV94 commented Jul 17, 2024

		compare_pytorch_and_py(x_fg, [np.r_[1.0, 2.0, 3.0, 4.0].astype(config.floatX)])


		def test_pytorch_Reshape_shape_graph_input():

	def test_pytorch_Reshape_shape_graph_input():
	def test_pytorch_Reshape_dynamic():

		x = DeepCopyOp()(pt.as_tensor_variable(1.1))
		x_fg = FunctionGraph([], [x])

Implements shape Ops and MakeVector in PyTorch #926

Implements shape Ops and MakeVector in PyTorch #926

Conversation

twaclaw commented Jul 12, 2024

Description

Related Issue

Checklist

Type of change

codecov bot commented Jul 12, 2024 • edited Loading

Codecov Report

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

Ch0ronomato Jul 16, 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

twaclaw Jul 16, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ricardoV94 Jul 17, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ricardoV94 commented Jul 17, 2024

codecov bot commented Jul 12, 2024 •

edited

Loading

Ch0ronomato Jul 16, 2024 •

edited

Loading

twaclaw Jul 16, 2024 •

edited

Loading

ricardoV94 Jul 17, 2024 •

edited

Loading