Update `std_mean/var_mean/nanmean/nansum` signatures with `int[1]? dim` (#82912)
### Description
Change the type of the `dim` arg for `std_mean/var_mean/nanmean/nansum` to `int[1]?` in `native_functions.yaml`
### Issue
Part of #29137
### Testing
Pull Request resolved: https://github.com/pytorch/pytorch/pull/82912
Approved by: https://github.com/albanD
diff --git a/aten/src/ATen/native/ReduceOps.cpp b/aten/src/ATen/native/ReduceOps.cpp
index 57f9eff..71fb7d9 100644
--- a/aten/src/ATen/native/ReduceOps.cpp
+++ b/aten/src/ATen/native/ReduceOps.cpp
@@ -1088,7 +1088,7 @@
return at::sum_out(result, self, dimnames_to_positions(self, dim), keepdim, opt_dtype);
}
-Tensor& nansum_out(const Tensor& self, IntArrayRef dim,
+Tensor& nansum_out(const Tensor& self, at::OptionalIntArrayRef dim,
bool keepdim, optional<ScalarType> opt_dtype, Tensor& result) {
TORCH_CHECK(!c10::isComplexType(self.scalar_type()), "nansum does not support complex inputs");
// For integral types, use existing sum as
@@ -1107,7 +1107,7 @@
return result;
}
-Tensor nansum(const Tensor& self, IntArrayRef dim, bool keepdim, c10::optional<ScalarType> opt_dtype) {
+Tensor nansum(const Tensor& self, at::OptionalIntArrayRef dim, bool keepdim, c10::optional<ScalarType> opt_dtype) {
ScalarType dtype = get_dtype_from_self(self, opt_dtype, true);
Tensor result = create_reduction_result(self, dim, keepdim, dtype);
return at::native::nansum_out(self, dim, keepdim, dtype, result);
@@ -1239,7 +1239,7 @@
// TODO(@heitorschueroff) implement custom kernels for nanmean
Tensor& nanmean_out(
const Tensor& self,
- IntArrayRef dim,
+ at::OptionalIntArrayRef dim,
bool keepdim,
c10::optional<ScalarType> opt_dtype,
Tensor& result) {
@@ -1254,7 +1254,7 @@
Tensor nanmean(
const Tensor& self,
- IntArrayRef dim,
+ at::OptionalIntArrayRef dim,
bool keepdim,
optional<ScalarType> opt_dtype) {
TORCH_CHECK(
@@ -1729,25 +1729,31 @@
}
std::tuple<Tensor, Tensor> var_mean(
- const Tensor& self, IntArrayRef dim, bool unbiased, bool keepdim) {
- return at::var_mean(self, /*dim=*/at::OptionalIntArrayRef(dim),
- /*correction=*/int64_t{unbiased ? 1 : 0}, keepdim);
+ const Tensor& self, at::OptionalIntArrayRef dim, bool unbiased, bool keepdim) {
+ return at::var_mean(
+ self, /*dim=*/at::OptionalIntArrayRef(dim),
+ /*correction=*/c10::make_optional<int64_t>({unbiased ? 1 : 0}),
+ keepdim);
}
std::tuple<Tensor, Tensor> std_mean(
- const Tensor& self, IntArrayRef dim, bool unbiased, bool keepdim) {
- return at::std_mean(self, /*dim=*/at::OptionalIntArrayRef(dim),
- /*correction=*/int64_t{unbiased ? 1 : 0}, keepdim);
+ const Tensor& self, at::OptionalIntArrayRef dim, bool unbiased, bool keepdim) {
+ return at::std_mean(
+ self, /*dim=*/at::OptionalIntArrayRef(dim),
+ /*correction=*/c10::make_optional<int64_t>({unbiased ? 1 : 0}),
+ keepdim);
}
std::tuple<Tensor, Tensor> std_mean(const Tensor& self, bool unbiased) {
return at::std_mean(
- self, /*dim=*/c10::nullopt, /*correction=*/int64_t{unbiased ? 1 : 0});
+ self, /*dim=*/c10::nullopt,
+ /*correction=*/c10::make_optional<int64_t>({unbiased ? 1 : 0}));
}
std::tuple<Tensor, Tensor> var_mean(const Tensor& self, bool unbiased) {
return at::var_mean(
- self, /*dim=*/c10::nullopt, /*correction=*/int64_t{unbiased ? 1 : 0});
+ self, /*dim=*/c10::nullopt,
+ /*correction=*/c10::make_optional<int64_t>({unbiased ? 1 : 0}));
}
std::tuple<Tensor&, Tensor&> var_mean_out(
diff --git a/aten/src/ATen/native/ReduceOpsUtils.h b/aten/src/ATen/native/ReduceOpsUtils.h
index 7c73c85..9db9802 100644
--- a/aten/src/ATen/native/ReduceOpsUtils.h
+++ b/aten/src/ATen/native/ReduceOpsUtils.h
@@ -159,7 +159,7 @@
}
inline Tensor create_reduction_result(
- const Tensor& self, IntArrayRef dim, bool keepdim, ScalarType dtype
+ const Tensor& self, at::OptionalIntArrayRef dim, bool keepdim, ScalarType dtype
) {
DimMask mask = make_dim_mask(dim, self.dim());
auto shape = shape_from_dim_mask(self, mask, keepdim);
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
index a7193d4..524e71c 100644
--- a/aten/src/ATen/native/native_functions.yaml
+++ b/aten/src/ATen/native/native_functions.yaml
@@ -3311,11 +3311,11 @@
- func: mean.names_out(Tensor self, Dimname[1] dim, bool keepdim=False, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!)
device_check: NoCheck # TensorIterator
-- func: nanmean(Tensor self, int[1] dim=[], bool keepdim=False, *, ScalarType? dtype=None) -> Tensor
+- func: nanmean(Tensor self, int[1]? dim=None, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor
device_check: NoCheck # Composite
variants: function, method
-- func: nanmean.out(Tensor self, int[1] dim=[], bool keepdim=False, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!)
+- func: nanmean.out(Tensor self, int[1]? dim=None, bool keepdim=False, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!)
device_check: NoCheck # Composite
- func: median(Tensor self) -> Tensor
@@ -4772,12 +4772,12 @@
- func: sum.DimnameList_out(Tensor self, Dimname[1] dim, bool keepdim=False, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!)
device_check: NoCheck # TensorIterator
-- func: nansum(Tensor self, int[1] dim=[], bool keepdim=False, *, ScalarType? dtype=None) -> Tensor
+- func: nansum(Tensor self, int[1]? dim=None, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor
variants: function, method
dispatch:
CPU, CUDA: nansum
-- func: nansum.out(Tensor self, int[1] dim=[], bool keepdim=False, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!)
+- func: nansum.out(Tensor self, int[1]? dim=None, bool keepdim=False, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!)
dispatch:
CPU, CUDA: nansum_out
@@ -4842,7 +4842,7 @@
device_check: NoCheck # TensorIterator
variants: function
-- func: std_mean.dim(Tensor self, int[1] dim, bool unbiased=True, bool keepdim=False) -> (Tensor, Tensor)
+- func: std_mean.dim(Tensor self, int[1]? dim, bool unbiased=True, bool keepdim=False) -> (Tensor, Tensor)
device_check: NoCheck # TensorIterator
variants: function
@@ -5279,7 +5279,7 @@
device_check: NoCheck # TensorIterator
variants: function
-- func: var_mean.dim(Tensor self, int[1] dim, bool unbiased=True, bool keepdim=False) -> (Tensor, Tensor)
+- func: var_mean.dim(Tensor self, int[1]? dim, bool unbiased=True, bool keepdim=False) -> (Tensor, Tensor)
device_check: NoCheck # TensorIterator
variants: function
diff --git a/test/forward_backward_compatibility/check_forward_backward_compatibility.py b/test/forward_backward_compatibility/check_forward_backward_compatibility.py
index 7f91b95..9a3d450 100644
--- a/test/forward_backward_compatibility/check_forward_backward_compatibility.py
+++ b/test/forward_backward_compatibility/check_forward_backward_compatibility.py
@@ -115,6 +115,10 @@
("aten::.*functional", datetime.date(2022, 8, 1)),
("aten::_foreach.*", datetime.date(2022, 8, 1)),
("aten::unflatten", datetime.date(2022, 8, 10)),
+ ("aten::nanmean", datetime.date(2022, 8, 30)),
+ ("aten::nanmean.out", datetime.date(2022, 8, 30)),
+ ("aten::nansum", datetime.date(2022, 8, 30)),
+ ("aten::nansum.out", datetime.date(2022, 8, 30)),
# TODO: FIXME: prims shouldn't be checked
("prims::.*", datetime.date(9999, 1, 1)),
]
diff --git a/test/test_namedtensor.py b/test/test_namedtensor.py
index d7ca5d6..751a56f 100644
--- a/test/test_namedtensor.py
+++ b/test/test_namedtensor.py
@@ -1181,7 +1181,17 @@
check_output(op(t, 1), ['N', 'L'])
check_output(op(t, -1), ['N', 'C'])
check_output(op(t, 'C'), ['N', 'L'])
- if op.__name__ in ['sum', 'mean', 'std', 'var']:
+ ops_support_dim_none = [
+ 'sum',
+ 'mean',
+ 'std',
+ 'var',
+ 'std_mean',
+ 'var_mean',
+ 'nanmean',
+ 'nansum',
+ ]
+ if op.__name__ in ops_support_dim_none:
check_output(op(t, None), [])
else:
with self.assertRaisesRegex(RuntimeError, 'Please look up dimensions by name'):
diff --git a/tools/autograd/derivatives.yaml b/tools/autograd/derivatives.yaml
index 4fccfd1..f3daf6e 100644
--- a/tools/autograd/derivatives.yaml
+++ b/tools/autograd/derivatives.yaml
@@ -1539,7 +1539,7 @@
self: sum_backward(grad, self.sizes(), dim, keepdim)
result: auto_linear
-- name: nansum(Tensor self, int[1] dim=[], bool keepdim=False, *, ScalarType? dtype=None) -> Tensor
+- name: nansum(Tensor self, int[1]? dim=None, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor
self: nansum_backward(grad.to(self.scalar_type()), self, dim, keepdim)
result: at::where(self_p.isnan(), 0, self_t).sum(dim, keepdim, dtype)
diff --git a/torch/_torch_docs.py b/torch/_torch_docs.py
index c552ad0..b699ea6 100644
--- a/torch/_torch_docs.py
+++ b/torch/_torch_docs.py
@@ -7095,7 +7095,7 @@
Args:
{input}
- {dim} If `None`, reduces all dimensions. Default is `None`.
+ {opt_dim}
{keepdim}
Keyword args:
@@ -10806,7 +10806,7 @@
Args:
{input}
- {dim}
+ {opt_dim}
Keyword args:
unbiased (bool): whether to use Bessel's correction (:math:`\delta N = 1`).
@@ -10971,7 +10971,7 @@
Args:
{input}
- {dim}
+ {opt_dim}
{keepdim}
Keyword args:
@@ -12228,7 +12228,7 @@
Args:
{input}
- {dim}
+ {opt_dim}
Keyword args:
unbiased (bool): whether to use Bessel's correction (:math:`\delta N = 1`).
diff --git a/torch/csrc/autograd/FunctionsManual.cpp b/torch/csrc/autograd/FunctionsManual.cpp
index 8934549..ccbed38 100644
--- a/torch/csrc/autograd/FunctionsManual.cpp
+++ b/torch/csrc/autograd/FunctionsManual.cpp
@@ -613,7 +613,7 @@
Tensor nansum_backward(
const Tensor& grad,
const Tensor& self,
- IntArrayRef dims,
+ at::OptionalIntArrayRef dims,
bool keepdim) {
return sum_backward(grad, self.sizes(), dims, keepdim) *
self.isnan().logical_not();
diff --git a/torch/csrc/autograd/FunctionsManual.h b/torch/csrc/autograd/FunctionsManual.h
index 376490f..8ae8641 100644
--- a/torch/csrc/autograd/FunctionsManual.h
+++ b/torch/csrc/autograd/FunctionsManual.h
@@ -165,7 +165,7 @@
at::Tensor nansum_backward(
const at::Tensor& grad,
const at::Tensor& self,
- at::IntArrayRef dims,
+ at::OptionalIntArrayRef dims,
bool keepdim);
std::vector<int64_t> reverse_list(const at::IntArrayRef list);
at::Tensor reverse_dim(const at::Tensor& t, int64_t dim);
