Add memory format support to `randn_like` operator (#27890)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/27890
Adds memory_format keyword argument (positional for cpp).
'Preserve' behavior now follows next rules:
1) If tensor is non-overlapping and dense - output tensor will have the same strides as input tensor.
2) If not (1) and tensor is stored in the channels last format, output tensor going to have channels last format.
3) Output tensor is going to be contiguous in all other cases.
---
Dense tensor is the tensor that store values in a contiguous block of memory.
Non-overlapping tensor is the tensor in which elements occupy individual non-repetitive memory.
Test Plan: Imported from OSS
Differential Revision: D17980314
Pulled By: VitalyFedyunin
fbshipit-source-id: a2cf3b1b2df1a4956da971fd47ce69487b2c09e9
diff --git a/aten/src/ATen/native/TensorFactories.cpp b/aten/src/ATen/native/TensorFactories.cpp
index 9191b04..2903d93 100644
--- a/aten/src/ATen/native/TensorFactories.cpp
+++ b/aten/src/ATen/native/TensorFactories.cpp
@@ -566,12 +566,18 @@
return result.normal_(mean, std, generator);
}
-Tensor randn_like(const Tensor& self) {
- return native::randn_like(self, self.options());
+Tensor randn_like(
+ const Tensor& self,
+ c10::optional<c10::MemoryFormat> optional_memory_format) {
+ return native::randn_like(self, self.options(), optional_memory_format);
}
-Tensor randn_like(const Tensor& self, const TensorOptions& options) {
- return native::randn(self.sizes(), nullptr, options);
+Tensor randn_like(
+ const Tensor& self,
+ const TensorOptions& options,
+ c10::optional<c10::MemoryFormat> optional_memory_format) {
+ auto result = at::empty_like(self, options, optional_memory_format);
+ return result.normal_(0, 1, nullptr);
}
// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ randperm ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
index c9f7720..67880f5 100644
--- a/aten/src/ATen/native/native_functions.yaml
+++ b/aten/src/ATen/native/native_functions.yaml
@@ -2105,10 +2105,9 @@
- func: randn.generator_out(int[] size, *, Generator? generator, Tensor(a!) out) -> Tensor(a!)
-- func: randn_like(Tensor self) -> Tensor
- use_c10_dispatcher: full
+- func: randn_like(Tensor self, *, MemoryFormat? memory_format=None) -> Tensor
-- func: randn_like.dtype(Tensor self, *, ScalarType dtype, Layout layout, Device device, bool pin_memory=False) -> Tensor
+- func: randn_like.dtype(Tensor self, *, ScalarType dtype, Layout layout, Device device, bool pin_memory=False, MemoryFormat? memory_format=None) -> Tensor
- func: randperm(int n, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
diff --git a/test/test_torch.py b/test/test_torch.py
index 91ec27c..6dd38f7 100644
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -12602,7 +12602,8 @@
lambda t, **kwargs: torch.zeros_like(t, **kwargs),
lambda t, **kwargs: torch.ones_like(t, **kwargs),
lambda t, **kwargs: torch.randint_like(t, 10, 100, **kwargs),
- lambda t, **kwargs: torch.randint_like(t, 100, **kwargs)]
+ lambda t, **kwargs: torch.randint_like(t, 100, **kwargs),
+ lambda t, **kwargs: torch.randn_like(t, **kwargs)]
for transformation_fn in transformation_fns:
self._test_memory_format_transformations(device, input_generator_fn, transformation_fn, compare_data=False)
diff --git a/tools/autograd/gen_variable_factories.py b/tools/autograd/gen_variable_factories.py
index 2d78aaf..194b512 100644
--- a/tools/autograd/gen_variable_factories.py
+++ b/tools/autograd/gen_variable_factories.py
@@ -66,7 +66,8 @@
SUPPORT_MEMORY_FORMAT = [
'empty_like', 'full_like',
'ones_like', 'rand_like',
- 'zeros_like', 'randint_like']
+ 'zeros_like', 'randint_like',
+ 'randn_like']
if decl['name'] in SUPPORT_MEMORY_FORMAT:
actuals.insert(-1, '{}.options().is_variable(false)'.format(actuals[0]))
else:
diff --git a/torch/csrc/jit/ir.cpp b/torch/csrc/jit/ir.cpp
index d88caa4..88d199e 100644
--- a/torch/csrc/jit/ir.cpp
+++ b/torch/csrc/jit/ir.cpp
@@ -909,8 +909,8 @@
"aten::randint_like(Tensor self, int high, *, int dtype, int layout, Device device, bool pin_memory, MemoryFormat? memory_format=None) -> Tensor",
"aten::randint_like(Tensor self, int low, int high, *, int dtype, int layout, Device device, bool pin_memory, MemoryFormat? memory_format=None) -> Tensor",
"aten::randn(int[] size, *, int? dtype, int? layout, Device? device, bool? pin_memory) -> Tensor",
- "aten::randn_like(Tensor self) -> Tensor",
- "aten::randn_like(Tensor self, *, int dtype, int layout, Device device, bool pin_memory) -> Tensor",
+ "aten::randn_like(Tensor self, *, MemoryFormat? memory_format=None) -> Tensor",
+ "aten::randn_like(Tensor self, *, int dtype, int layout, Device device, bool pin_memory, MemoryFormat? memory_format=None) -> Tensor",
"aten::randperm(int n, *, int? dtype, int? layout, Device? device, bool? pin_memory) -> Tensor"};
if (nondeterministic_ops.find(this) == nullptr) {
diff --git a/torch/csrc/jit/passes/shape_analysis.cpp b/torch/csrc/jit/passes/shape_analysis.cpp
index d704c1f..6d192f9 100644
--- a/torch/csrc/jit/passes/shape_analysis.cpp
+++ b/torch/csrc/jit/passes/shape_analysis.cpp
@@ -862,7 +862,7 @@
"aten::rand_like(Tensor self, *, MemoryFormat? memory_format=None) -> Tensor",
"aten::randint_like(Tensor self, int high, *, MemoryFormat? memory_format=None) -> Tensor",
"aten::randint_like(Tensor self, int low, int high, *, MemoryFormat? memory_format=None) -> Tensor",
- "aten::randn_like(Tensor self) -> Tensor",
+ "aten::randn_like(Tensor self, *, MemoryFormat? memory_format=None) -> Tensor",
"aten::zeros_like(Tensor self, *, MemoryFormat? memory_format=contiguous_format) -> Tensor",
},
[](Node* node) -> type_vec_t {
@@ -1417,7 +1417,7 @@
"aten::rand_like(Tensor self, *, int dtype, int layout, Device device, bool pin_memory, MemoryFormat? memory_format=contiguous_format) -> Tensor",
"aten::randint_like(Tensor self, int high, *, int dtype, int layout, Device device, bool pin_memory, MemoryFormat? memory_format=contiguous_format) -> Tensor",
"aten::randint_like(Tensor self, int low, int high, *, int dtype, int layout, Device device, bool pin_memory, MemoryFormat? memory_format=contiguous_format) -> Tensor",
- "aten::randn_like(Tensor self, *, int dtype, int layout, Device device, bool pin_memory) -> Tensor",
+ "aten::randn_like(Tensor self, *, int dtype, int layout, Device device, bool pin_memory, MemoryFormat? memory_format=None) -> Tensor",
"aten::zeros_like(Tensor self, *, int dtype, int layout, Device device, bool pin_memory, MemoryFormat? memory_format=contiguous_format) -> Tensor",
},
[](Node* node) -> type_vec_t {
