Rewrite functional.tensordot to be TorchScript-able #53672

gmagogsfm · 2021-03-10T01:09:41Z

Fixes #53487

facebook-github-bot · 2021-03-10T01:09:51Z

💊 CI failures summary and remediations

As of commit d8ec82c (more details on the Dr. CI page):

2/2 failures possibly* introduced in this PR
- 2/2 non-scanned failure(s)

ci.pytorch.org: 1 failed

Failed: pr/pytorch-linux-bionic-rocm4.0.1-py3.6

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions to the (internal) Dr. CI Users group.

facebook-github-bot

@gmagogsfm has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

tugsbayasgalan · 2021-03-10T19:40:04Z

torch/functional.py

    r"""Returns a contraction of a and b over multiple dimensions.

    :attr:`tensordot` implements a generalized matrix product.

    Args:
      a (Tensor): Left tensor to contract
      b (Tensor): Right tensor to contract
-      dims (int or Tuple[List[int]] containing two lists): number of dimensions to
+      dims (int or Tuple[List[int], List[int]] or List[List[int]] containing two lists or Tensor): number of dimensions to


In the official documentation (https://pytorch.org/docs/stable/generated/torch.tensordot.html), it looks like we don't support List[int] or List[List[int]]. Or is it necessary because of how aten op is defined?

The documentation link you referred to is actually generated from this docstring, so when I update it in this PR, documentation will change.

The support for List[int] is already there in the implementation, I am updating docstring to match what's implemented.

eellison

I dont think we should be adding overhead to eager opperators to support them in JIT. This function can be supported with overloads while keeping the original code more in tact.

eellison · 2021-03-10T19:58:27Z

torch/functional.py

-    if isinstance(dims, (list, tuple)) or \
-       (isinstance(dims, torch.Tensor) and dims.numel() > 1):
+
+    dims_a = torch.jit.annotate(List[int], [])


this isnt up to date syntax, you should be using dims_a : List[int] = []

eellison · 2021-03-10T20:01:33Z

torch/functional.py

@@ -937,15 +937,15 @@ def _consecutive_return_inverse(input, return_inverse=False, return_counts=False
 unique_consecutive.__doc__ = _unique_consecutive_impl.__doc__


-def tensordot(a, b, dims=2, out=None):
+def tensordot(a, b, dims: Any = 2, out: Optional[torch.Tensor] = None):


torch.jit.isinstance and adding overhead here, we're iterating over the entire list when we use to just check if the input was a list type. I dont think we should be adding overhead to eager operators.

I think you can support this without making these changes by using overloads:

@overload
def tensordot(a, b, dims: int = 2, out: Optional[torch.Tensor] = None):

@overload
def tensordot(a, b, dims: Tuple[List[int], List[int]], out: Optional[torch.Tensor] = None):

@overload
def tensordot(a, b, dims: List[List[int]], out: Optional[torch.Tensor] = None):

Good call! Changed to use @overload.

gmagogsfm

Revised, PTAL

facebook-github-bot

@gmagogsfm has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

facebook-github-bot

@gmagogsfm has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

eellison

Looks good, thanks for doing this ! some of the assertions may be better explicit as you have them, idk, just commenting to let you know where they will be checked as part of the following line

eellison · 2021-03-11T20:07:22Z

torch/functional.py

-        if isinstance(dims, torch.Tensor):
-            dims = dims.item()
+
+    if isinstance(dims, list):


Nit you can fold this if with the tuple one.
if isinstance(dims, (list, tuple))

eellison · 2021-03-11T20:07:55Z

torch/functional.py

-            dims = dims.item()
+
+    if isinstance(dims, list):
+        assert len(dims) == 2


Nit the assertion here is unnecessary since the list unpacking will also check that the length is 2

eellison · 2021-03-11T20:08:44Z

torch/functional.py

+        num_elements = dims.numel()
+        if num_elements > 1:
+            assert dims.size()[0] == 2
+            dims_a = torch.jit.annotate(List[int], dims[0].tolist())


Nit: not to use torch.jit.annotate use up to date syntax

Unfortunately, this doesn't work at the moment. TS complains about redefining variable dims_a.

file an issue and add in comment ?

eellison · 2021-03-11T20:25:28Z

torch/functional.py

+            dims_a = torch.jit.annotate(List[int], dims[0].tolist())
+            dims_b = torch.jit.annotate(List[int], dims[1].tolist())
+        else:
+            assert num_elements == 1


nit: dims.item() will already check that len == 1

eellison · 2021-03-11T20:26:43Z

torch/functional.py

+            dims_b = torch.jit.annotate(List[int], dims[1].tolist())
+        else:
+            assert num_elements == 1
+            dims_val = int(dims.item())


I think you're missing a check that dims_val >= 0 here

Good catch, added.

facebook-github-bot

@gmagogsfm has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

codecov · 2021-03-12T00:26:32Z

Codecov Report

Merging #53672 (d8ec82c) into master (0584fd9) will decrease coverage by 0.00%.
The diff coverage is 78.78%.

@@            Coverage Diff             @@
##           master   #53672      +/-   ##
==========================================
- Coverage   77.36%   77.35%   -0.01%     
==========================================
  Files        1879     1879              
  Lines      183257   183285      +28     
==========================================
+ Hits       141774   141785      +11     
- Misses      41483    41500      +17

facebook-github-bot · 2021-03-21T06:05:18Z

@gmagogsfm merged this pull request in f48a971.

fritzo · 2021-07-01T02:59:00Z

torch/functional.py

+    if len(dims_a) == 0 or len(dims_b) == 0:
+        raise RuntimeError(f"unsupported input to tensordot, got dims={dims}")


This check breaks the simple case tensordot(torch.tensor(0.), torch.tensor(0.), 0), which works fine in NumPy

>>> np.tensordot(np.zeros(()), np.zeros(()), 0) array(0.)

This important edge case dims=0 is needed by generic machinery as in Pyro and opt_einsum. Can we revert this check?

cc @neerajprad

facebook-github-bot added oncall: jit Add this issue/PR to JIT oncall triage queue cla signed labels Mar 10, 2021

gmagogsfm requested review from eellison and tugsbayasgalan March 10, 2021 02:12

facebook-github-bot reviewed Mar 10, 2021

View reviewed changes

tugsbayasgalan approved these changes Mar 10, 2021

View reviewed changes

eellison requested changes Mar 10, 2021

View reviewed changes

gmagogsfm commented Mar 11, 2021

View reviewed changes

gmagogsfm requested a review from eellison March 11, 2021 02:00

facebook-github-bot reviewed Mar 11, 2021

View reviewed changes

eellison approved these changes Mar 11, 2021

View reviewed changes

Rewrite functional.tensordot to be TorchScript-able

d8ec82c

facebook-github-bot reviewed Mar 11, 2021

View reviewed changes

facebook-github-bot closed this in f48a971 Mar 21, 2021

facebook-github-bot added the Merged label Mar 21, 2021

This was referenced Apr 18, 2021

Make tensordot resize output tensor's size if out= argument is specified & make it safely cast & copy output #56286

Closed

test_variant_consistency_jit fails for torch.tensordot with dtype float32 with error INTERNAL ASSERT FAILED #56314

Open

fritzo reviewed Jul 1, 2021

View reviewed changes

fritzo mentioned this pull request Jul 1, 2021

torch.tensordot(-,-,0) no longer works #61096

Closed

		if len(dims_a) == 0 or len(dims_b) == 0:
		raise RuntimeError(f"unsupported input to tensordot, got dims={dims}")

Rewrite functional.tensordot to be TorchScript-able #53672

Rewrite functional.tensordot to be TorchScript-able #53672

Uh oh!

Conversation

gmagogsfm commented Mar 10, 2021

Uh oh!

facebook-github-bot commented Mar 10, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

💊 CI failures summary and remediations

ci.pytorch.org: 1 failed

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

eellison left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gmagogsfm left a comment

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

eellison left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

codecov bot commented Mar 12, 2021

Codecov Report

Uh oh!

facebook-github-bot commented Mar 21, 2021

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

facebook-github-bot commented Mar 10, 2021 •

edited

Loading