Merge pull request #1101 from dwf/squared_l2_norm

Bug fix + enhancement to l2_norm.

blocks/theano_expressions.py

@@ -2,7 +2,7 @@
 from theano import tensor
 
 
-def l2_norm(tensors):
+def l2_norm(tensors, squared=False):
     """Computes the total L2 norm of a set of tensors.
 
     Converts all operands to :class:`~tensor.TensorVariable`
@@ -12,11 +12,13 @@ def l2_norm(tensors):
     ----------
     tensors : iterable of :class:`~tensor.TensorVariable` (or compatible)
         The tensors.
+    squared : bool, optional
+        If `True`, return the squared L2 norm. Default: `False`.
 
     """
     summed = [tensor.sqr(tensor.as_tensor_variable(t)).sum() for t in tensors]
-    joined = tensor.stack(*summed)
-    return tensor.sqrt(joined.sum())
+    joined = tensor.stack(summed, axis=0)
+    return joined.sum() if squared else tensor.sqrt(joined.sum())
 
 
 def hessian_times_vector(gradient, parameter, vector, r_op=False):

tests/test_theano_expressions.py

@@ -12,6 +12,8 @@ def test_l2_norm():
     assert_allclose(l2_norm([3, [1, 2]]).eval(), 14.0 ** 0.5)
     assert_allclose(
         l2_norm([3, [1, 2], [[1, 2], [3, 4]]]).eval(), 44.0 ** 0.5)
+    assert_allclose(
+        l2_norm([3, [1, 2], [[1, 2], [3, 4]]], squared=True).eval(), 44.0)
 
 
 def test_hessian_times_vector():