test_resize_as()
+ def test_requires_grad_loop(self):
+ @torch.jit.script
+ def test(x, y, z):
+ # type: (Tensor, Tensor, int) -> Tensor
+ for _ in range(z):
+ x = y
+ return x
+
+ # x requires grad, y does not
+ # testing that requires grad analysis correctly exits, with its input
+ # to the loop (x) requiring grad and its output to the loop not requiring grad
+ # and the output of the node conservatively setting grad to true
+
+ inps = (torch.tensor(1.0, requires_grad=True), torch.tensor(1), 10)
+ test(*inps)
+
+ graph = test.graph_for(*inps)
+ loop = graph.findNode("prim::Loop")
+ loop_body = next(loop.blocks())
+ loop_inputs = list(loop_body.inputs())
+ loop_outputs = list(loop_body.outputs())
+
+ self.assertTrue(loop_inputs[1].requires_grad())
+ self.assertFalse(loop_outputs[1].requires_grad())
+ self.assertTrue(loop.output().requires_grad())
+
def test_view_shape_prop(self):
cu = torch.jit.CompilationUnit('''
def test_view_shape_prop(a):
+#include <ATen/core/jit_type.h>
#include <torch/csrc/jit/argument_spec.h>
#include <torch/csrc/jit/ir.h>
#include <torch/csrc/jit/operator.h>
-#include <ATen/core/jit_type.h>
#include <vector>
fmap(node->inputs().slice(2), getRequiresGrad);
std::vector<bool> body_outputs_require(node->outputs().size(), false);
- while (body_inputs_require != body_outputs_require) {
- body_inputs_require =
+ std::vector<bool> new_body_inputs_require = body_inputs_require;
+ std::vector<bool> new_body_outputs_require = body_outputs_require;
+
+ // continue iterating until the results have converged
+ do {
+ body_inputs_require = new_body_inputs_require;
+ body_outputs_require = new_body_outputs_require;
+
+ new_body_inputs_require =
bitwiseOr(body_inputs_require, body_outputs_require);
setRequiresGrad(
- body->param_node()->outputs().slice(1), body_inputs_require);
+ body->param_node()->outputs().slice(1), new_body_inputs_require);
PropagateRequiresGrad(body);
- body_outputs_require =
+ new_body_outputs_require =
fmap(body->return_node()->inputs().slice(1), getRequiresGrad);
- }
+ } while (new_body_inputs_require != body_inputs_require &&
+ new_body_outputs_require != body_outputs_require);
setRequiresGrad(node, body_outputs_require);
} else {
PropagateRequiresGrad(node);
}
}
-
} // anonymous namespace
void PropagateRequiresGrad(std::shared_ptr<Graph>& graph) {
PropagateRequiresGrad(graph->block());
}
-
} // namespace jit
} // namespace torch