From 80f4374dde4342625b22a4f415bfb92fdc1290d2 Mon Sep 17 00:00:00 2001 From: Sebastian Messmer Date: Mon, 28 Jan 2019 11:36:30 -0800 Subject: [PATCH] Handle stack correctly (#16246) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/16246 The op schema says it returns multiple values, so let's actually return multiple values instead of one tuple. For some reason, this did work when called from python (probably some auto-unpacking), but once called from JIT, it segfaulted. This diff fixes that. Reviewed By: dzhulgakov Differential Revision: D13780147 fbshipit-source-id: fe94f82f4c53b7454f77c4484fca4ac9dc444475 --- caffe2/python/operator_test/layer_norm_op_test.py | 21 +++++++++++++++++++++ torch/csrc/jit/c10_ops/layer_norm.cpp | 4 ++-- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/caffe2/python/operator_test/layer_norm_op_test.py b/caffe2/python/operator_test/layer_norm_op_test.py index 892b9cc..3f8fa75 100644 --- a/caffe2/python/operator_test/layer_norm_op_test.py +++ b/caffe2/python/operator_test/layer_norm_op_test.py @@ -133,6 +133,27 @@ class TestLayerNormOp(serial.SerializedTestCase): torch.testing.assert_allclose(expected_stdev, actual_stdev) @given(X=hu.tensors(n=1), **hu.gcs) + def test_layer_norm_op_jit(self, X, gc, dc): + @torch.jit.script + def jit_layer_norm(tensor, axis, epsilon): + # type: (Tensor, int, float) -> Tuple[Tensor, Tensor, Tensor] + norm, mean, stdev = torch.ops.caffe2.layer_norm_dont_use_this_op_yet(tensor, axis, epsilon) + return norm, mean, stdev + + X = X[0] + if len(X.shape) == 1: + X = np.expand_dims(X, axis=0) + axis = np.random.randint(0, len(X.shape)) + epsilon = 1e-4 + + actual_norm, actual_mean, actual_stdev = jit_layer_norm(torch.tensor(X), axis, epsilon) + expected_norm, expected_mean, expected_stdev = _layer_norm_ref(axis, epsilon, X) + + torch.testing.assert_allclose(expected_norm, actual_norm) + torch.testing.assert_allclose(expected_mean, actual_mean) + torch.testing.assert_allclose(expected_stdev, actual_stdev) + + @given(X=hu.tensors(n=1), **hu.gcs) def test_layer_norm_brew_wrapper(self, X, gc, dc): X = X[0] if len(X.shape) == 1: diff --git a/torch/csrc/jit/c10_ops/layer_norm.cpp b/torch/csrc/jit/c10_ops/layer_norm.cpp index 30366ca..d1b3809 100644 --- a/torch/csrc/jit/c10_ops/layer_norm.cpp +++ b/torch/csrc/jit/c10_ops/layer_norm.cpp @@ -54,11 +54,11 @@ RegisterOperators reg({ epsilon, IValue(cache) }); - push(stack, Tuple::create({ + push(stack, torch::autograd::make_variable(at::Tensor(std::move(c10_output)), false), torch::autograd::make_variable(at::Tensor(std::move(c10_output_mean)), false), torch::autograd::make_variable(at::Tensor(std::move(c10_output_stdev)), false) - })); + ); return 0; }) }); -- 2.7.4