From 2c190d2f05a538dba8762d071c69b20945fa0004 Mon Sep 17 00:00:00 2001
From: James Cross <jcross@fb.com>
Date: Fri, 1 Dec 2017 19:59:20 -0800
Subject: [PATCH] update transformer code for layer_norm() API change

Summary: Quick fix for unit test broken by D6454290. This is my fault for approving while the tests covering the single callsite were broken.

Reviewed By: goldsborough

Differential Revision: D6466566

fbshipit-source-id: 2683be3d6bb184286e64fbde3e572946e39030c7
---
 caffe2/python/helpers/normalization.py        | 20 +++++++++++++------
 .../operator_test/layer_norm_op_test.py       |  2 +-
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/caffe2/python/helpers/normalization.py b/caffe2/python/helpers/normalization.py
index 7845fe5d42b..168d785fbec 100644
--- a/caffe2/python/helpers/normalization.py
+++ b/caffe2/python/helpers/normalization.py
@@ -216,29 +216,37 @@ def layer_norm(
     # The learned multiplicative scale or "gain".
     scale = model.create_param(
         param_name='{}_scale'.format(blob_out),
-        shape=dim_in,
-        initializer=initializers.Initializer('ConstantFill', value=initial_scale),
-        tags=ParameterTags.WEIGHT
+        shape=[dim_in],
+        initializer=initializers.Initializer(
+            'ConstantFill',
+            value=initial_scale,
+        ),
+        tags=ParameterTags.WEIGHT,
     )
 
     # The learned additive bias or "shift".
     bias = model.create_param(
         param_name='{}_bias'.format(blob_out),
-        shape=dim_in,
-        initializer=initializers.Initializer('ConstantFill', value=initial_bias),
-        tags=ParameterTags.BIAS
+        shape=[dim_in],
+        initializer=initializers.Initializer(
+            'ConstantFill',
+            value=initial_bias,
+        ),
+        tags=ParameterTags.BIAS,
     )
 
     scaled = model.net.Mul(
         [normalized, scale],
         ['{}_scaled'.format(blob_out)],
         broadcast=1,
+        axis=axis,
     )
 
     biased = model.net.Add(
         [scaled, bias],
         ['{}_biased'.format(blob_out)],
         broadcast=1,
+        axis=axis,
     )
 
     return biased, mean, stdev
diff --git a/caffe2/python/operator_test/layer_norm_op_test.py b/caffe2/python/operator_test/layer_norm_op_test.py
index ba7d5159c38..2b443f1b3c0 100644
--- a/caffe2/python/operator_test/layer_norm_op_test.py
+++ b/caffe2/python/operator_test/layer_norm_op_test.py
@@ -155,7 +155,7 @@ class TestLayerNormOp(hu.HypothesisTestCase):
             model,
             'input',
             'output',
-            dim_in=scale_dim,
+            dim_in=X.shape[axis],
             axis=axis,
             epsilon=1e-4,
         )