def test_c10_layer_norm(self, X, gc, dc): X = X[0] if len(X.shape) == 1: X = np.expand_dims(X, axis=0) axis = np.random.randint(0, len(X.shape)) scale_dim = [1] * np.ndim(X) scale_dim[axis] = X.shape[axis] self.ws.create_blob('input').feed(X) model = ModelHelper(name='test_layer_norm_brew_wrapper') model.C10LayerNorm_DontUseThisOpYet( ['input'], ['output_1', 'output_1', 'output_2'], axis=axis, epsilon=1e-4, ) self.ws.create_net(model.param_init_net).run() self.ws.create_net(model.net).run()