PyTorch Softmax Ops #846

HAKSOAT · 2024-06-23T22:31:27Z

Description

This PR implements the Softmax Ops for PyTorch.

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

…o pytensor-pytorch

…ensor-pytorch-softmax

…ftmax

HAKSOAT · 2024-06-23T22:40:56Z

Hi @ricardoV94,

I went through this PR: #764 and I observed that at some point in that work, Softmax, LogSoftmax and SoftmaxGrad were added, in this commit.

So I'm curious why it was taken out. I see something about cuda, but was hoping to get more context.

ricardoV94 · 2024-06-23T22:43:47Z

Hi @ricardoV94,

I went through this PR: #764 and I observed that at some point in that work, Softmax, LogSoftmax and SoftmaxGrad were added, in this commit.

So I'm curious why it was taken out. I see something about cuda, but was hoping to get more context.

We tried to reduce the scope of that initial PR to get the basics in and merge it sooner rather than later. Dropping the implementation of those Ops was probably just that

ricardoV94 · 2024-06-23T22:44:59Z

pytensor/link/pytorch/dispatch/elemwise.py

+def pytorch_funcify_Softmax(op, **kwargs):
+    axis = op.axis
+
+    if axis is None:


Hmm, None means all axis in PyTensor, doesn't Pytorch support that case?

None in Pytorch implicitly means axis 1: https://discuss.pytorch.org/t/implicit-dimension-choice-for-softmax-warning/12314/10

And that funcationality is deprecated now, giving the warning:

UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument.
torch.nn.functional.softmax(torch.Tensor(x), dim=None)

So can we write the function to achieve the same behavior as softmax along all axis? I guess in that case we need to do it manually with exp(x) / torch.sum(exp(x), axis=None) (with the max subtraction for stability) or is that also not allowed?

We shouldn't raise an error that is related to the torch API, but try to achieve the intended behavior of PyTensor. Whether we use torch.softmax or something else is an implementation detail for PyTensor users.

Or something like: if axis is None: torch.softmax(x.ravel()).reshape(x.shape)?

Yes, reshaping can be used for axis=None
Also, it might be a good idea to make sure x is a torch.float dtype. The softmax would fail if the dtype is int or long.

@ricardoV94 , what is the expected behaviour when the input is of dtype int?

I'm thinking of converting to float implicitly.

Have to check the integer case but don't forget to test it as well

Hi @ricardoV94, I have added the Logsoftmax and Softmaxgrad ops.

I have a question in relation to testing the types. How do I do this?

Everytime I send the functiongraph and the inputs into compare_pytorch_and_py, the function always ends up with a Tensor of floats somehow.

So can you help with pointers on what other way I can test the actual inputs without some conversion happening in the middle. I am asking because of the need to check that the Ops still work with an input of dtype int.

In the test you're creating a matrix which has a default dtype of float32 or float64, you can parametrize the test with an explicit dtype, say pytest.mark.parametrize("dtype", ["float64", "int64"]) to test integer types as well, and create x = matrix("x", dtype=dtype)

This is something we should mention in the docs @HarshvirSandhu

ricardoV94 · 2024-06-24T06:53:32Z

CC @HarshvirSandhu for helping with the review as the PR progresses

HAKSOAT · 2024-06-25T23:35:41Z

environment.yml

@@ -9,7 +9,7 @@ channels:
 dependencies:
  - python>=3.10
  - compilers
-  - numpy>=1.17.0
+  - numpy>=1.17.0,<2


This change is because of this error: #827

ricardoV94 · 2024-06-26T08:21:33Z

pytensor/link/pytorch/dispatch/elemwise.py

+
+    def softmax(x):
+        if not torch.is_floating_point(x):
+            x = x.to(torch.float32)


Should probably convert to the output type advertised by PyTensor. The second argument of the funcify function, is node (now hidden inside **kwargs). If you retrieve node, you can check the dtype of the output via node.outputs[0].dtype. This is probably what you should convert to, not necessarily float32

Suggested change

x = x.to(torch.float32)

x = x.to(torch.float32)

Actually there's a bug in PyTensor Softmax, it also fails if you try to execute with integer dtype in the default backend. I'll open an issue for that.

For now it's enough for the torch dispatch function to raise a NotImplementedError if the input dtype (which similarly you can get from node.inputs[0].dtype is an integer, and not try to handle it inside the returned softmax function.

Same for the other Softmax related functions

Opened an issue: #857

ricardoV94 · 2024-06-27T09:13:25Z

pytensor/link/pytorch/dispatch/elemwise.py

@@ -40,11 +40,14 @@ def dimshuffle(x):
 @pytorch_funcify.register(Softmax)
 def pytorch_funcify_Softmax(op, **kwargs):
    axis = op.axis
+    dtype = kwargs["node"].outputs[0].dtype


Maybe better to check the input dtype, because we would fail if we pass an integer. PyTensor could start saying Softmax takes as input integers and outputs floats once we fix it? Sorry if I said the output before

Suggested change

dtype = kwargs["node"].outputs[0].dtype

dtype = kwargs["node"].inputs[0].dtype

ricardoV94 · 2024-06-27T09:14:37Z

Left one comment, after that I think it's ready!

…/pytensor into pytensor-pytorch-softmax

…ensor-pytorch-softmax

HAKSOAT · 2024-06-28T10:20:21Z

Hi @ricardoV94 , I have made that change and resolved existing merge conflicts.

codecov · 2024-06-28T11:12:11Z

Codecov Report

Attention: Patch coverage is 73.33333% with 8 lines in your changes missing coverage. Please review.

Project coverage is 80.98%. Comparing base (920b409) to head (b4cdce0).
Report is 148 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/link/pytorch/dispatch/elemwise.py	73.33%	8 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #846      +/-   ##
==========================================
- Coverage   80.98%   80.98%   -0.01%     
==========================================
  Files         169      169              
  Lines       46985    47025      +40     
  Branches    11494    11501       +7     
==========================================
+ Hits        38052    38084      +32     
- Misses       6716     6727      +11     
+ Partials     2217     2214       -3

Files with missing lines	Coverage Δ
pytensor/link/pytorch/dispatch/elemwise.py	`69.81% <73.33%> (+4.59%)`	⬆️

... and 3 files with indirect coverage changes

ricardoV94 · 2024-06-28T11:14:38Z

Thanks @HAKSOAT

HarshvirSandhu added 30 commits May 13, 2024 16:01

Add pytorch support for some basic Ops

27e2526

update variable names, docstrings

629d00b

Avoid numpy conversion of torch Tensors

3eceb56

Fix typify and CheckAndRaise

3cde964

Fix Elemwise Ops

c003aa5

Fix Scalar Ops

8dc406e

Fix ruff-format

a8f6ddb

Initial setup for pytorch tests

9d535f5

Fix mode parameters for pytorch

c5600da

Prevent conversion of scalars to numpy

54b6248

Update TensorConstantSignature and map dtypes to Tensor types

19454b3

Add tests for basic ops

92d7114

Remove torch from user facing API

5aae0e5

Add function to convert numpy arrays to pytorch tensors

8c174dd

Avoid copy when converting to tensor

0977c3a

Fix tests

1c23825

Remove dispatches that are not tested

c9195a8

set path for pytorch tests

b07805c

Remove tensorflow probability from yml

9e8d3fc

Add checks for runtime broadcasting

a2d3afa

Remove IfElse

a577a80

Remove dev notebook

499a174

Fix check and raise

2826613

Fix compare_pytorch_and_py

62ffcec

Fix DimShuffle

acdbba1

Add tests for Elemwise operations

2519c65

Fix test for CheckAndRaise

eb6d5c2

Remove duplicate function

9f02a4f

Remove device from pytorch_typify

caf2965

Merge branch 'main' of https://github.com/HarshvirSandhu/pytensor int…

bf87eb9

…o pytensor-pytorch

HarshvirSandhu and others added 6 commits June 17, 2024 17:10

remove device from elemwise tests and add assertions

8ec7661

skip tests if cuda is not available

bb7df41

Fix tests

0441cf2

Merge branch 'main' of https://github.com/pymc-devs/pytensor into pyt…

85f2742

…ensor-pytorch-softmax

Implemented softmax ops for PyTorch

4ca5aca

Merge remote-tracking branch 'upstream/main' into pytensor-pytorch-so…

b9aca57

…ftmax

ricardoV94 reviewed Jun 23, 2024

View reviewed changes

ricardoV94 mentioned this pull request Jun 24, 2024

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

Open

48 tasks

HAKSOAT added 2 commits June 24, 2024 21:56

Switched to run softmax on all items if axis is None

287d9c2

Implemented log softmax

f42e2a0

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

Implemented softmaxgrad

35b17e0

HAKSOAT commented Jun 25, 2024

View reviewed changes

ricardoV94 reviewed Jun 26, 2024

View reviewed changes

ricardoV94 mentioned this pull request Jun 26, 2024

Softmax fails with integer dtypes only at runtime #857

Open

Added checks and error raises for nonfloat inputs

5efc3c8

ricardoV94 reviewed Jun 27, 2024

View reviewed changes

HAKSOAT added 3 commits June 28, 2024 11:12

Added checks and error raises for nonfloat inputs

16e415a

Merge branch 'pytensor-pytorch-softmax' of https://github.com/HAKSOAT…

ffbc594

…/pytensor into pytensor-pytorch-softmax

Merge branch 'main' of https://github.com/pymc-devs/pytensor into pyt…

b4cdce0

…ensor-pytorch-softmax

ricardoV94 approved these changes Jun 28, 2024

View reviewed changes

ricardoV94 merged commit 17fa8b1 into pymc-devs:main Jun 28, 2024
56 of 57 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

PyTorch Softmax Ops #846

PyTorch Softmax Ops #846

HAKSOAT commented Jun 23, 2024

HAKSOAT commented Jun 23, 2024

ricardoV94 commented Jun 23, 2024

ricardoV94 Jun 23, 2024

HAKSOAT Jun 23, 2024

ricardoV94 Jun 24, 2024 •

edited

Loading

ricardoV94 Jun 24, 2024

HarshvirSandhu Jun 24, 2024

HAKSOAT Jun 24, 2024

ricardoV94 Jun 25, 2024

HAKSOAT Jun 25, 2024 •

edited

Loading

ricardoV94 Jun 26, 2024 •

edited

Loading

ricardoV94 commented Jun 24, 2024

HAKSOAT Jun 25, 2024

ricardoV94 Jun 26, 2024 •

edited

Loading

ricardoV94 Jun 26, 2024 •

edited

Loading

ricardoV94 Jun 26, 2024

ricardoV94 Jun 27, 2024 •

edited

Loading

ricardoV94 commented Jun 27, 2024

HAKSOAT commented Jun 28, 2024

codecov bot commented Jun 28, 2024 •

edited

Loading

ricardoV94 commented Jun 28, 2024

	dtype = kwargs["node"].outputs[0].dtype
	dtype = kwargs["node"].inputs[0].dtype

PyTorch Softmax Ops #846

PyTorch Softmax Ops #846

Conversation

HAKSOAT commented Jun 23, 2024

Description

Related Issue

Checklist

Type of change

HAKSOAT commented Jun 23, 2024

ricardoV94 commented Jun 23, 2024

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

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

HAKSOAT Jun 25, 2024 • edited Loading

Choose a reason for hiding this comment

ricardoV94 Jun 26, 2024 • edited Loading

Choose a reason for hiding this comment

ricardoV94 commented Jun 24, 2024

Choose a reason for hiding this comment

ricardoV94 Jun 26, 2024 • edited Loading

Choose a reason for hiding this comment

ricardoV94 Jun 26, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ricardoV94 Jun 27, 2024 • edited Loading

Choose a reason for hiding this comment

ricardoV94 commented Jun 27, 2024

HAKSOAT commented Jun 28, 2024

codecov bot commented Jun 28, 2024 • edited Loading

Codecov Report

ricardoV94 commented Jun 28, 2024

ricardoV94 Jun 24, 2024 •

edited

Loading

HAKSOAT Jun 25, 2024 •

edited

Loading

ricardoV94 Jun 26, 2024 •

edited

Loading

ricardoV94 Jun 26, 2024 •

edited

Loading

ricardoV94 Jun 26, 2024 •

edited

Loading

ricardoV94 Jun 27, 2024 •

edited

Loading

codecov bot commented Jun 28, 2024 •

edited

Loading