update comment and tests

theano/tensor/nnet/bn.py

@@ -52,9 +52,11 @@ def batch_normalization(inputs, gamma, beta, mean, std,
     mode: 'low_mem' or 'high_mem'
         Specify which batch_normalization implementation that will be
         used.
-        As no intermediate representations are stored for the
-        back-propagation, 'low_mem' implementation lower the memory usage, however,
-        it is 5-10% slower than 'high_mem' implementation.
+        As no intermediate representations are stored for the back-propagation,
+        'low_mem' implementation lower the memory usage, however,
+        it is 5-10% slower than 'high_mem' implementation. Note that 5-10% computation
+        time difference compare the batch_normalization operation only, time difference
+        between implementation is likely to be less important on the full model fprop/bprop.
     """
     if mode == 'low_mem':
         elm_bn = theano.tensor.elemwise.Elemwise(scalar_op=BNComposite(dtype=inputs.dtype))

theano/tensor/nnet/tests/test_bn.py

@@ -41,7 +41,7 @@ def test_bn():
     f_ref = theano.function([x, b, g], [bn_ref_op])
     res_ref = f_ref(X, G, B)
     for mode in ['low_mem', 'high_mem']:
-        bn_op = batch_normalization(x, g, b, x.mean(axis=0, keepdims=True), x.std(axis=0, keepdims=True))
+        bn_op = batch_normalization(x, g, b, x.mean(axis=0, keepdims=True), x.std(axis=0, keepdims=True), mode=mode)
         f = theano.function([x, b, g], [bn_op])
         res = f(X, G, B)
         utt.assert_allclose(res_ref, res)