8000 Implement `np.diff` for single order differences by soulitzer · Pull Request #50569 · pytorch/pytorch · GitHub
[go: up one dir, main page]

Skip to content

Implement np.diff for single order differences #50569

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 15 commits into from
Closed
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Only handle tensor append and prepend
  • Loading branch information
soulitzer committed Jan 25, 2021
commit 19bb4d461038a5356828d1df5daa76d8ca85ec8a
39 changes: 2 additions & 37 deletion 8000 s aten/src/ATen/native/ReduceOps.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -448,7 +448,7 @@ static inline Tensor diff_helper(const Tensor& self, int64_t n, int64_t dim) {
return at::narrow(self, dim, 1, out_len) - at::narrow(self, dim, 0, out_len);
}

Tensor diff_tensor_tensor(const Tensor& self, int64_t n, int64_t dim, const c10::optional<Tensor>& prepend, const c10::optional<Tensor>& append) {
Tensor diff(const Tensor& self, int64_t n, int64_t dim, const c10::optional<Tensor>& prepend, const c10::optional<Tensor>& append) {
if (!prepend.has_value() && !append.has_value()) {
return diff_helper(self, n, dim);
} else {
Expand All @@ -465,7 +465,7 @@ static inline Tensor& diff_out_helper(const Tensor& self, int64_t n, int64_t dim
return at::sub_out(result, at::narrow(self, dim, 1, out_len), at::narrow(self, dim, 0, out_len));
}

Tensor& diff_tensor_tensor_out(const Tensor& self, int64_t n, int64_t dim, const c10::optional<Tensor>& prepend, const c10::optional<Tensor>& append, Tensor& result) {
Tensor& diff_out(const Tensor& self, int64_t n, int64_t dim, const c10::optional<Tensor>& prepend, const c10::optional<Tensor>& append, Tensor& result) {
if (!prepend.has_value() && !append.has_value()) {
return diff_out_helper(self, n, dim, result);
} else {
Expand All @@ -474,41 +474,6 @@ Tensor& diff_tensor_tensor_out(const Tensor& self, int64_t n, int64_t dim, const
}
}

static Tensor diff_broadcast_scalar(Scalar scalar, const Tensor& tensor, int64_t dim) {
// Helper for diff to handle when prepend/append are scalars
return at::scalar_tensor(scalar, tensor.options()).broadcast_to(tensor.sizes()).narrow(dim, 0, 1);
}

Tensor diff_scalar_scalar(const Tensor& self, int64_t n, int64_t dim, c10::optional<Scalar> prepend, c10::optional<Scalar> append) {
return diff_tensor_tensor(self, n, dim,
prepend.has_value() ? c10::optional<Tensor>(diff_broadcast_scalar(prepend.value(), self, dim)) : c10::nullopt,
append.has_value() ? c10::optional<Tensor>(diff_broadcast_scalar(append.value(), self, dim)) : c10::nullopt);
}

Tensor diff_scalar_tensor(const Tensor& self, int64_t n, int64_t dim, c10::optional<Scalar> prepend, const c10::optional<Tensor>& append) {
return diff_tensor_tensor(self, n, dim, prepend.has_value() ? c10::optional<Tensor>(diff_broadcast_scalar(prepend.value(), self, dim)) : c10::nullopt, append);
}

Tensor diff_tensor_scalar(const Tensor& self, int64_t n, int64_t dim, const c10::optional<Tensor>& prepend, c10::optional<Scalar> append) {
return diff_tensor_tensor(self, n, dim, prepend, append.has_value() ? c10::optional<Tensor>(diff_broadcast_scalar(append.value(), self, dim)) : c10::nullopt);
}

Tensor& diff_scalar_scalar_out(const Tensor& self, int64_t n, int64_t dim, c10::optional<Scalar> prepend, c10::optional<Scalar> append, Tensor& result) {
return diff_tensor_tensor_out(self, n, dim,
prepend.has_value() ? c10::optional<Tensor>(diff_broadcast_scalar(prepend.value(), self, dim)) : c10::nullopt,
append.has_value() ? c10::optional<Tensor>(diff_broadcast_scalar(append.value(), self, dim)) : c10::nullopt, result);
}

Tensor& diff_scalar_tensor_out(const Tensor& self, int64_t n, int64_t dim, c10::optional<Scalar> prepend, const c10::optional<Tensor>& append, Tensor& result) {
return diff_tensor_tensor_out(self, n, dim,
prepend.has_value() ? c10::optional<Tensor>(diff_broadcast_scalar(prepend.value(), self, dim)) : c10::nullopt, append, result);
}

Tensor& diff_tensor_scalar_out(const Tensor& self, int64_t n, int64_t dim, const c10::optional<Tensor>& prepend, c10::optional<Scalar> append, Tensor& result) {
return diff_tensor_tensor_out(self, n, dim, prepend,
append.has_value() ? c10::optional<Tensor>(diff_broadcast_scalar(append.value(), self, dim)) : c10::nullopt, result);
}

// ALL REDUCE #################################################################

static ScalarType get_dtype(Tensor& result, const Tensor& self, optional<ScalarType> dtype,
Expand Down
44 changes: 4 additions & 40 deletions aten/src/ATen/native/native_functions.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -1373,51 +1373,15 @@
- func: fill_diagonal_(Tensor(a!) self, Scalar fill_value, bool wrap=False) -> Tensor(a!)
variants: method

- func: diff.Tensor_Tensor(Tensor self, int n=1, int dim=-1, Tensor? prepend=None, Tensor? append=None) -> Tensor
- func: diff(Tensor self, int n=1, int dim=-1, Tensor? prepend=None, Tensor? append=None) -> Tensor
Copy link
Collaborator
@mruberry mruberry Jan 27, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cc @rgommers for a possible issue with NumPy's signature

np.diff has both prepend and append default to "no value", which is represented by:

https://github.com/numpy/numpy/blob/1de46fe2feee9d3c500a83331ac9b75af5aef947/numpy/_globals.py#L57

The comment indicates it's intended to be used with "deprecated keywords." However, it appears these defaults come from the PR adding the kwargs: numpy/numpy#8206. @rgommers, seems like the NumPy signature would ideally have prepend=None and append=None?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That comment is a little misleading, or at least incomplete. _NoValue is mainly used when keywords are added to a function, and the function forwards its keywords somewhere. See for example np.sum which forwards to the .sum method of its input (which can be a torch.Tensor). Adding keywords to the numpy function that are not yet present in all other/downstream libraries implementing the same functionality would then break previously working code. I'll go fix the code comment.

Wherever you see _NoValue, just read None.

The PR gives this example that would not work with prepend=None for object arrays:

>>> data = np.array([2, 3], dtype=object)
>>> np.diff(data, prepend=None)
array([None, 1], dtype=object)  # expected
array([1], dtype=object)  # actual

A little far-fetched, but unfortunately this works:

>>> np.array(None)                                                          
array(None, dtype=object)

Object arrays are very annoying.

variants: function, method
dispatch:
Math: diff_tensor_tensor
Math: diff

- func: diff.Tensor_Scalar(Tensor self, int n=1, int dim=-1, Tensor? prepend=None, Scalar? append=None) -> Tensor
cpp_no_default_args: ['n', 'dim', 'prepend', 'append']
variants: function, method
dispatch:
Math: diff_tensor_scalar

- func: diff.Scalar_Tensor(Tensor self, int n=1, int dim=-1, Scalar? prepend=None, Tensor? append=None) -> Tensor
cpp_no_default_args: ['n', 'dim', 'prepend', 'append']
variants: function, method
dispatch:
Math: diff_scalar_tensor

- func: diff.Scalar_Scalar(Tensor self, int n=1, int dim=-1, Scalar? prepend=None, Scalar? append=None) -> Tensor
cpp_no_default_args: ['n', 'dim', 'prepend', 'append']
variants: function, method
dispatch:
Math: diff_scalar_scalar

- func: diff.Tensor_Tensor_out(Tensor self, int n=1, int dim=-1, Tensor? prepend=None, Tensor? append=None, *, Tensor(a!) out) -> Tensor(a!)
variants: function
dispatch:
Math: diff_tensor_tensor_out

- func: diff.Tensor_Scalar_out(Tensor self, int n=1, int dim=-1, Tensor? prepend=None, Scalar? append=None, *, Tensor(a!) out) -> Tensor(a!)
cpp_no_default_args: ['n', 'dim', 'prepend', 'append']
variants: function
dispatch:
Math: diff_tensor_scalar_out

- func: diff.Scalar_Tensor_out(Tensor self, int n=1, int dim=-1, Scalar? prepend=None, Tensor? append=None, *, Tensor(a!) out) -> Tensor(a!)
cpp_no_default_args: ['n', 'dim', 'prepend', 'append']
variants: function
dispatch:
Math: diff_scalar_tensor_out

- func: diff.Scalar_Scalar_out(Tensor self, int n=1, int dim=-1, Scalar? prepend=None, Scalar? append=None, *, Tensor(a!) out) -> Tensor(a!)
cpp_no_default_args: ['n', 'dim', 'prepend', 'append']
- func: diff.out(Tensor self, int n=1, int dim=-1, Tensor? prepend=None, Tensor? append=None, *, Tensor(a!) out) -> Tensor(a!)
variants: function
dispatch:
Math: diff_scalar_scalar_out
Math: diff_out

- func: div.Tensor(Tensor self, Tensor other) -> Tensor
variants: function, method
Expand Down
39 changes: 14 additions & 25 deletions test/test_torch.py
Original file line number Diff line number Diff line change
Expand Up @@ -4106,7 +4106,10 @@ def logcumsumexp(a, axis):
'expected scalar_type Double but found Float'):
torch.logcumsumexp(b, axis, out=inplace_out)

def test_diff(self, device):
@dtypes(*torch.testing.get_all_dtypes(include_bfloat16=False, include_half=False))
@dtypesIfCPU(*torch.testing.get_all_dtypes(include_bfloat16=False, include_half=True))
@dtypesIfCUDA(*torch.testing.get_all_dtypes(include_bfloat16=True, include_half=True))
def test_diff(self, device, dtype):
def _test_diff_numpy(t, dims=None):
for dim in dims if dims else range(t.dim()):
actual = torch.diff(t, dim=dim, prepend=t, append=t)
Expand All @@ -4120,43 +4123,29 @@ def _test_diff_numpy(t, dims=None):
expected = torch.from_numpy(np.diff(np_t, axis=dim, prepend=np_t, append=np_t))
self.assertEqual(actual, expected.to(t.dtype))

scalar = t[(0,) * t.dim()]
scalar_np = np_t[(0,) * t.dim()]

# test when both prepend and append are scalars
actual = torch.diff(t, dim=dim, prepend=scalar, append=scalar)
expected = torch.from_numpy(np.diff(np_t, axis=dim, prepend=scalar_np, append=scalar_np))
self.assertEqual(actual, expected.to(t.dtype))

# test when prepend is a scalar and append is a tensor
actual = torch.diff(t, dim=dim, prepend=scalar, append=t)
expected = torch.from_numpy(np.diff(np_t, axis=dim, prepend=scalar_np, append=np_t))
self.assertEqual(actual, expected.to(t.dtype))

shapes = (
(1,),
(1, 5),
(3, 5),
(1, 5, 1),
(2, 3, 5))

for dt in torch.testing.get_all_dtypes():
for shape in shapes:
contig = _make_tensor(shape, dt, device)
_test_diff_numpy(contig)
for shape in shapes:
contig = _make_tensor(shape, dtype, device)
_test_diff_numpy(contig)

non_contig = torch.empty(shape + (2, 2), dtype=dt)[..., 0]
non_contig = non_contig.select(-1, -1)
non_contig.copy_(contig)
self.assertTrue(not non_contig.is_contiguous() or shape == (1,))
non_contig = torch.empty(shape + (2, 2), dtype=dtype)[..., 0]
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Use make_tensor instead of _make_tensor:

1E79
def make_tensor(size, device: torch.device, dtype: torch.dtype, *,

Both here on 4152 and on 4149.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've replaced _make_tensor with make_tensor on line 4149, but I don't think we can use make_tensor to produce torch.empty for line 4152.

non_contig = non_contig.select(-1, -1)
non_contig.copy_(contig)
self.assertTrue(not non_contig.is_contiguous() or shape == (1,))
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cool testing of both contiguous and discontiguous use cases.


_test_diff_numpy(non_contig)
_test_diff_numpy(non_contig)

t = torch.ones(2, 3)

# Size of prepend/append needs to be same as t except along given dim
with self.assertRaises(RuntimeError):
invalid_prepend = torch.tensor([[0, 1]])
invalid_prepend = torch.tensor([[0, 1]], device=device, dtype=dtype)
t.diff(dim=0, prepend=invalid_prepend)

with self.assertRaisesRegex(
Expand All @@ -4165,7 +4154,7 @@ def _test_diff_numpy(t, dims=None):

with self.assertRaisesRegex(
RuntimeError, 'diff requires input that is at least one-dimensional'):
scalar = torch.tensor(2)
scalar = torch.tensor(2, device=device, dtype=dtype)
torch.diff(scalar)

def _test_large_cum_fn_helper(self, x, fn):
Expand Down
10 changes: 4 additions & 6 deletions torch/_torch_docs.py
Original file line number Diff line number Diff line change
Expand Up @@ -2651,12 +2651,10 @@ def merge_dicts(*dicts):
n (int, optional): the number of times to recursively compute difference
dim (int, optional): the dimension with respect to compute the difference.
Default is the last dimension.
prepend, append (Tensor or Scalar, optional): values to prepend or append to
:attr:`input` along :attr:`dim` before computing the difference. Scalar
values are expanded to tensors with size 1 along :attr:`dim` and the shape of
:attr:`input` along all other dimensions. Otherwise, its dimension must be
equivalent to that of input, and its shape must match input's shape except on
:attr:`dim`.
prepend, append (Tensor, optional): values to prepend or append to
:attr:`input` along :attr:`dim` before computing the difference.
Its dimension must be equivalent to that of input, and its shape
must match input's shape except on :attr:`dim`.

Keyword args:
{out}
Expand Down
0