from __future__ import absolute_import, division, print_function

import pytest

import tensorflow as tf
from lucid.optvis.param.unit_balls import unit_ball_L2, unit_ball_L_inf


learning_rate = .1
num_steps = 16


@pytest.mark.skip(reason="Unpredictably imprecise right now.")
@pytest.mark.parametrize("shape", [(3), (2, 5, 5)])
def test_unit_ball_L2(shape, eps=1e-6):
  """Tests that a L2 unit ball variable's norm stays roughly within 1.0.
  Note: only holds down to eps ~= 5e-7.
  """
  with tf.Session() as sess:
    unit_ball = unit_ball_L2(shape)
    unit_ball_L2_norm = tf.norm(unit_ball)
    optimizer = tf.train.GradientDescentOptimizer(learning_rate)
    objective = optimizer.minimize(-unit_ball)
    tf.global_variables_initializer().run()
    norm_value = unit_ball_L2_norm.eval()
    for i in range(num_steps):
      _, new_norm_value = sess.run([objective, unit_ball_L2_norm])
      assert new_norm_value >= norm_value - eps
      assert new_norm_value <= 1.0 + eps
      norm_value = new_norm_value


@pytest.mark.skip(reason="Unpredictably imprecise right now.")
@pytest.mark.parametrize("shape", [(3), (2, 5, 5)])
@pytest.mark.parametrize("precondition", [True, False])
def test_unit_ball_L_inf(shape, precondition, eps=1e-6):
  """Tests that a L infinity unit ball variables' stay roughly within 1.0.
  Note: only holds down to eps ~= 5e-7.
  """
  with tf.Session() as sess:
    unit_ball = unit_ball_L_inf(shape, precondition=precondition)
    unit_ball_max = tf.reduce_max(unit_ball)
    optimizer = tf.train.GradientDescentOptimizer(learning_rate)
    objective = optimizer.minimize(-unit_ball)
    tf.global_variables_initializer().run()
    for i in range(num_steps):
      _, unit_ball_max_value = sess.run([objective, unit_ball_max])
      assert unit_ball_max_value <= 1.0 + eps

@@ -98,16 +98,31 @@
 w as tf%0A
+import logging%0A
 from luc
@@ -221,16 +221,50 @@
 ps = 16%0A
+log = logging.getLogger(__name__)%0A
 %0A%0A@pytes
@@ -262,36 +262,37 @@
 %0A%0A%0A@pytest.mark.
-skip
+xfail
 (reason=%22Unpredi
@@ -300,32 +300,42 @@
 tably imprecise 
+on Travis 
 right now.%22)%0A@py
@@ -898,24 +898,41 @@
 w_norm_value
+, unit_ball_value
  = sess.run(
@@ -960,17 +960,97 @@
 _L2_norm
-%5D
+, unit_ball%5D)%0A      log.info(%22L2: %25s, Value: %25s%22, new_norm_value, unit_ball_value
 )%0A      
@@ -1185,12 +1185,13 @@
 ark.
-skip
+xfail
 (rea
@@ -1219,16 +1219,26 @@
 precise 
+on Travis 
 right no
@@ -1879,24 +1879,41 @@
 ll_max_value
+, unit_ball_value
  = sess.run(
@@ -1937,17 +1937,104 @@
 ball_max
-%5D
+, unit_ball%5D)%0A      log.info(%22Linf: %25s, Value: %25s%22, unit_ball_max_value, unit_ball_value
 )%0A      

Mark unit_ball tests as xfail. Large imprecision observed on CI can not be reproduced locally. Adds some logging to help investigate.