[primTorch] slice and transpose & etc. #76727

mruberry · 2022-05-03T08:35:41Z

This PR...

Adds the following prims:

slice
slice_in_dim
transpose

Adds the following refs:

cat
permute
transpose
swap_axes (alias for transpose)
tensor_split

Makes the following test improvements:

adds reference inputs for torch.permute
adds a NumPy reference for torch.permute
adds reference inputs for torch.cat

Fixes the following bugs:

adds support for scalars to the min and max prims

facebook-github-bot · 2022-05-03T08:35:48Z

🔗 Helpful links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/76727
📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓Need help or want to give feedback on the CI? Visit our office hours
↩️ [fb-only] Re-run with SSH instructions

💊 CI failures summary and remediations

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

Expand to see more

✅ None of the CI failures appear to be your fault 💚

2/2 broken upstream at merge base 9b97313 since May 03

🚧 2 ongoing upstream failures:

These were probably caused by upstream breakages that are not fixed yet.

pull / linux-bionic-rocm5.1-py3.7 / test (default, 1, 2, linux.rocm.gpu) since May 03 (c51b53d)
- 🔁 rerun
pull / pytorch-xla-linux-bionic-py3.7-clang8 / test (xla, 1, 1, linux.2xlarge) since May 03 (9b97313)
- 🔁 rerun

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

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

Click here to manually regenerate this comment.

ngimel · 2022-05-03T23:29:56Z

torch/_prims/__init__.py

+    a: Union[TensorLikeType, NumberType], b: Union[TensorLikeType, NumberType]
+) -> TensorLikeType:
+    if isinstance(a, TensorLike) and isinstance(b, Number):
+        b = utils.wrap_scalar(b, dtype=a.dtype)


this looks wrong for cuda inputs (they accept scalar tensors on cpu, but here device checks on prim would error out)

This happens after those checks, though

ngimel · 2022-05-03T23:33:03Z

torch/_prims/__init__.py

+        )
+        raise ValueError(msg)
+
+    if a.ndim != len(start_indices):


it's the same condition as in line 916, so you've errored out already?

what is our general stance on where error checking like this should happen? E.g. for reductions I do it in ref, and prim assumes dims it receives is valid. You've already canonicalized indices before sending here, hence no negative dims, doesn't it follow that you've also error-checked them?

Correct, this should be limit_indices here

I think the prim should do its own error checking and in the future we should create an error context so the name of the called function appears in the error message, and try to minimize redundant error checking in the ref

(from a trace perspective only prim error checking matters)

ngimel · 2022-05-03T23:42:01Z

torch/_prims/__init__.py

+            raise ValueError(msg)
+        if x > y:
+            msg = (
+                "Attempting to slice a tensor but a stop index in {0} is greater than the length of "


is it stop index or end index?

In the Python slice docs it's start, stop and step: https://docs.python.org/3/c-api/slice.html

ngimel · 2022-05-03T23:43:33Z

torch/_prims/__init__.py

+            msg = (
+                "Attempting to slice a tensor but a stop index in {0} is greater than the length of "
+                " its corresponding dimension in shape {1}".format(
+                    limit_indices, a.shape


why full lists here unlike other error messages?

ngimel · 2022-05-03T23:45:34Z

torch/_prims/__init__.py

+
+    new_shape = []
+    for x, y in zip(start_indices, limit_indices):
+        new_shape.append(y - x)


should it be divided by stride?

yes! Good catch

ngimel · 2022-05-03T23:53:22Z

torch/_prims/__init__.py

+        raise ValueError(msg)
+
+    start_indices = [0] * a.ndim
+    limit_indices = [0] * a.ndim


should limit_indices be a.shape?

Yes! -- that's weird I thought I fixed that

ngimel · 2022-05-03T23:54:03Z

torch/_prims/__init__.py

+    axis: int = 0,
+) -> Tensor:
+    start_indices = [0] * a.ndim
+    limit_indices = list(a.shape)


yep they should be

ngimel · 2022-05-04T00:00:29Z

torch/_prims/utils.py

    """
-    return torch.tensor(a, dtype=type_to_dtype(type(a)))
+    if dtype is None:
+        return torch.tensor(a, dtype=type_to_dtype(type(a)))


It's confusing that now type promotion is split between ref (that converts tensor) and here (converts scalar). Maybe wrapping indeed should be done in ref.

Yeah I'm going to look at scalar handling more generally and ensure we're preserving scalars properly

ngimel · 2022-05-04T00:01:55Z

torch/_refs/__init__.py

+        msg = "cat expects at least one tensor, but received zero!"
+        raise ValueError(msg)
+
+    _dim = utils.canonicalize_dims(tensors[0].ndim, dim)


why not canonicalize_idx?

Yep, might as well

ngimel · 2022-05-04T00:08:21Z

torch/_refs/__init__.py

+
+    if out is not None:
+        out = _maybe_resize_out(out, result.shape)
+        return copy_to(out, result, allow_cross_device=False)  # type: ignore[arg-type]


Isn't it weird that all prims are not type-promoting except copy_to that is? Would it be better to instead check out dtype here, and not rely on copy_to to error out? If refs are documentation in code, then making error conditions clear is one of the goals.

Let's reconsider this when we apply out handling as a wrapper?

mruberry · 2022-05-04T05:36:55Z

@pytorchbot merge this please

Summary: This PR... Adds the following prims: - slice - slice_in_dim - transpose Adds the following refs: - cat - permute - transpose - swap_axes (alias for transpose) - tensor_split Makes the following test improvements: - adds reference inputs for torch.permute - adds a NumPy reference for torch.permute - adds reference inputs for torch.cat Fixes the following bugs: - adds support for scalars to the min and max prims Pull Request resolved: #76727 Approved by: https://github.com/ngimel Test Plan: contbuild & OSS CI, see https://hud.pytorch.org/commit/pytorch/pytorch/ef9f56eb0b86a2224652f5c5ac1f82dcb4f444c6 Reviewed By: malfet Differential Revision: D36134094 Pulled By: malfet fbshipit-source-id: 89e44eff92cf1b4bbc8ce123e45ee96ec580435e

shape ops

b2fa63f

mruberry requested a review from ngimel as a code owner May 3, 2022 08:35

facebook-github-bot added the cla signed label May 3, 2022

mruberry added topic: not user facing topic category module: primTorch labels May 3, 2022

mruberry changed the title ~~shape ops~~ [primTorch] slice and transpose & etc. May 3, 2022

Mike Ruberry added 3 commits May 3, 2022 01:39

removes temporary file

4389db9

excludes elu testing

9f3f94e

removes draft elu

d7ec2a9

ngimel reviewed May 4, 2022

View reviewed changes

Mike Ruberry added 2 commits May 3, 2022 20:04

stashes

4758125

merges and fixes per review

aefb50c

ngimel approved these changes May 4, 2022

View reviewed changes

Mike Ruberry added 2 commits May 3, 2022 20:39

stashes

1905406

Merge branch 'master' of ssh://github.com/pytorch/pytorch into shape_ops

be6a6e7

pytorchmergebot closed this in ef9f56e May 4, 2022

ngimel mentioned this pull request May 5, 2022

PrimTorch binary refs do not handle CUDA + CPU scalar tensors correctly #76871

Open

mruberry deleted the shape_ops branch May 19, 2022 18:59

[primTorch] slice and transpose & etc. #76727

[primTorch] slice and transpose & etc. #76727

Uh oh!

Conversation

mruberry commented May 3, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

facebook-github-bot commented May 3, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful links

💊 CI failures summary and remediations

🚧 2 ongoing upstream failures:

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!

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!

mruberry commented May 4, 2022

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

mruberry commented May 3, 2022 •

edited

Loading

facebook-github-bot commented May 3, 2022 •

edited

Loading