pytorch · bdhirsh · Apr 20, 2022 · Apr 20, 2022 · Apr 20, 2022 · Apr 21, 2022
diff --git a/aten/src/ATen/native/Fill.cpp b/aten/src/ATen/native/Fill.cpp
@@ -61,6 +61,14 @@ Tensor& fill_meta_(Tensor& self, const Tensor& value) {
   return self;
 }
 
+Tensor fill(const Tensor& self, const Scalar& value) {
+  return at::empty_like(self).fill_(value);
+}
+
+Tensor fill(const Tensor& self, const Tensor& value) {
+  return at::empty_like(self).fill_(value);
+}
+
 DEFINE_DISPATCH(fill_stub);
 
 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ fill_diagonal ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -2116,6 +2116,16 @@
 - func: unflatten.Dimname(Tensor(a) self, Dimname dim, int[] sizes, Dimname[] names) -> Tensor(a)
   variants: method
 
+- func: fill.Scalar(Tensor self, Scalar value) -> Tensor
+  variants: function
+  dispatch:
+    CompositeExplicitAutograd: fill
+
+- func: fill.Tensor(Tensor self, Tensor value) -> Tensor
+  variants: function
+  dispatch:
+    CompositeExplicitAutograd: fill
+
 - func: fill_.Scalar(Tensor(a!) self, Scalar value) -> Tensor(a!)
   device_check: NoCheck   # TensorIterator
   variants: function, method

diff --git a/test/jit/test_remove_mutation.py b/test/jit/test_remove_mutation.py
@@ -146,16 +146,16 @@ def test_successful():
 
         # full_like is not implemented for a tensor fill value
 
-        def test_unsuccessful():
+        def test_successful():
             x = torch.tensor([2, 2])
             y = torch.tensor([2, 4])
             x.fill_(y)
             return x + x
 
-        fn = torch.jit.script(test_unsuccessful)
+        fn = torch.jit.script(test_successful)
         graph = fn.graph
         self.run_pass('remove_mutation', graph)
-        FileCheck().check('aten::fill_').run(graph)
+        FileCheck().check_not('aten::fill_').run(graph)
 
         def normal():
             return torch.rand(2, 1, 3, 4).normal_()

diff --git a/test/test_functionalization.py b/test/test_functionalization.py
@@ -444,12 +444,28 @@ def f(x):
 
         self.assert_functionalization(f, torch.ones(2, 2))
         logs = self.get_logs(f, torch.ones(2, 2))
+        # zero() should decompose into zeros_like(), which will show up in the trace
         self.assertExpectedInline('\n'.join(logs), """\
 $0 = input('input')
 $1 = torch._ops.aten.add.Tensor($0, $0)
 $2 = torch._ops.aten.diagonal_copy.default($1)
 $3 = torch._ops.aten.zeros_like.default($2)""")
 
+    def test_fill_(self):
+        def f(x):
+            y = x + x
+            z = y.diagonal()
+            z.fill_(0)
+            return y
+
+        self.assert_functionalization(f, torch.ones(2, 2))
+        logs = self.get_logs(f, torch.ones(2, 2))
+        self.assertExpectedInline('\n'.join(logs), """\
+$0 = input('input')
+$1 = torch._ops.aten.add.Tensor($0, $0)
+$2 = torch._ops.aten.diagonal_copy.default($1)
+$3 = torch._ops.aten.fill.Scalar($2, 0)""")
+
     def test_nested_functions_propagate_updates(self):
         def g(x):
             # Create a view of x

diff --git a/tools/autograd/derivatives.yaml b/tools/autograd/derivatives.yaml
@@ -630,6 +630,15 @@
 - name: _fused_moving_avg_obs_fq_helper(Tensor self, Tensor observer_on, Tensor fake_quant_on, Tensor(a!) running_min, Tensor(b!) running_max, Tensor(c!) scale, Tensor(d!) zero_point, float averaging_const, int quant_min, int quant_max, int ch_axis, bool per_row_fake_quant=False, bool symmetric_quant=False) -> (Tensor output, Tensor mask)
   self: fake_quantize_per_tensor_affine_cachemask_backward(grad, mask)
 
+- name: fill.Scalar(Tensor self, Scalar value) -> Tensor
+  self: zeros_like(grad)
+  result: at::fill(self_t, 0)
+
+- name: fill.Tensor(Tensor self, Tensor value) -> Tensor
+  self: zeros_like(grad)
+  value: grad.sum()
+  result: at::fill(self_t, value_t)
+
 - name: fill_.Scalar(Tensor(a!) self, Scalar value) -> Tensor(a!)
   self: zeros_like(grad)
   result: self_t.fill_(0)

diff --git a/tools/autograd/gen_python_functions.py b/tools/autograd/gen_python_functions.py
@@ -153,6 +153,8 @@
     "replace_",  # only used by the functionalization pass, doesn't need to be exposed to python
     "zero",  # only used by the functionalization pass, doesn't need to be exposed to python
     "copy",  # only used by the functionalization pass
+    "fill.Tensor",  # only used by the functionalization pass
+    "fill.Scalar",  # only used by the functionalization pass
 ]
 
 SKIP_PYTHON_BINDINGS = list(

diff --git a/torch/overrides.py b/torch/overrides.py
@@ -143,6 +143,7 @@ def get_ignored_functions() -> Set[Callable]:
         torch.fft.rfftfreq,
         torch.from_file,
         torch.full,
+        torch.fill,
         torch.hamming_window,
         torch.hann_window,
         torch.kaiser_window,