Allow empty Variables to be saved for backwards (pytorch#23618)

Summary:
Pull Request resolved: pytorch#23618

For example: `save_for_backward({Variable(), x, Variable()})` should be allowed, so that this is consistent with the python API behaviour.

Test Plan: Added a test similar to the python test `test_save_none_for_backward` from test_autograd.py.

Differential Revision: D16589402

fbshipit-source-id: 847544ad8fc10772954d8629ad5a62bfdc1a66c1

test/cpp/api/autograd.cpp

@@ -127,3 +127,24 @@ TEST(CustomAutogradTest, ReturnLeafInplace) {
   q.sum().backward();
   ASSERT_VARIABLE_EQ(y.grad(), torch::ones({5,5}));
 }
+
+TEST(CustomAutogradTest, SaveEmptyForBackward) {
+  struct MyFunction : public Function<MyFunction> {
+    static variable_list forward(AutogradContext *ctx, Variable input) {
+      ctx->save_for_backward({Variable(), input, Variable()});
+      return {input*input};
+    }
+
+    static variable_list backward(AutogradContext *ctx, variable_list grad_output) {
+      auto saved = ctx->get_saved_variables();
+      EXPECT_FALSE(saved[0].defined());
+      EXPECT_FALSE(saved[2].defined());
+      return {saved[1] * 2 * grad_output[0]};
+    }
+  };
+
+  Variable x = torch::randn({5,5}, torch::requires_grad());
+  auto y = MyFunction::apply(x)[0];
+  y.sum().backward();
+  ASSERT_VARIABLE_EQ(x.grad(), 2*x);
+}

torch/csrc/autograd/custom_function.cpp

@@ -111,8 +111,13 @@ void AutogradContext::save_variables() {
   auto ptr = grad_fn_.lock();
 
   for (const auto& var : to_save_) {
-    bool is_output = var.grad_fn().get() == ptr.get();
-    saved_variables_.emplace_back(var, is_output);
+    // Allow empty variables to be saved
+    if (var.defined()) {
+      bool is_output = var.grad_fn().get() == ptr.get();
+      saved_variables_.emplace_back(var, is_output);
+    } else {
+      saved_variables_.emplace_back();
+    }
   }
   to_save_.clear();
 }