def test_bisect_batching(self):
"""Tests that bisect works in batching mode."""
wolfe_threshold = 1e-6
# Let's build our example function with 4 batches, each evaluating a
# different poly. They all have negative slopes both on 0.0 and 1.0,
# but different slopes (positive, negative) and values (low enough, too
# high) on their midpoint.
x = np.array([0.0, 0.5, 1.0])
y = np.array([[1.0, 0.6, 1.2],
[1.0, 0.6, 1.2],
[1.0, 1.6, 1.2],
[1.0, 1.6, 1.2]])
dy = np.array([[-0.8, 0.6, -0.7],
[-0.8, -0.4, -0.7],
[-0.8, 0.8, -0.7],
[-0.8, -0.4, -0.7]])
fun = test_function_x_y_dy(x, y, dy)
val_a = hzl._apply(fun, tf.zeros(4)) # Values at zero.
val_b = hzl._apply(fun, tf.ones(4)) # Values at initial step.
f_lim = val_a.f + (wolfe_threshold * tf.abs(val_a.f))
expected_left = np.array([0.0, 0.5, 0.0, 0.0])
expected_right = np.array([0.5, 0.75, 0.5, 0.25])
result = self.evaluate(hzl.bisect(fun, val_a, val_b, f_lim))
self.assertTrue(np.all(result.stopped))
self.assertTrue(np.all(~result.failed))
self.assertTrue(np.all(result.left.df < 0))
self.assertTrue(np.all(result.right.df >= 0))
self.assertArrayNear(result.left.x, expected_left, 1e-5)
self.assertArrayNear(result.right.x, expected_right, 1e-5)
def test_bisect_infty(self):
"""Tests that bisection is robust to +-inf, but fails on nan."""
wolfe_threshold = 1e-6
def fdf(x):
f = (x - 2.0)**2 # Minimum at 2
df = 2 * (x - 2.0)
# Cut the function to infinity at 4.
# Different "bad" values at different batch members.
infs = tf.constant([float('inf'), float('-inf'), float('nan')])
nans = tf.constant([float('nan'), float('nan'), float('nan')])
return ValueAndGradient(x=x,
f=tf.where(x > 4, infs, f),
df=tf.where(x > 4, nans, df))
val_a = fdf(tf.constant([0.0, 0.0, 0.0]))
val_b = fdf(tf.constant([5.0, 5.0, 5.0])) # Value isn't finite
f_lim = 4 + wolfe_threshold * 4
result = self.evaluate(hzl.bisect(fdf, val_a, val_b, f_lim))
self.assertTrue(np.all(result.stopped))
self.assertAllEqual(result.failed, [False, False, True])
self.assertAllEqual((result.left.df < 0), [True, False, True])
self.assertAllEqual((result.right.df >= 0), [True, False, False])
# TODO(axch): Actually expect something sensible
expected_left = np.array([0.0, 5.0, 0.0])
expected_right = np.array([2.5, 5.0, 5.0])
self.assertArrayNear(result.left.x, expected_left, 1e-5)
self.assertArrayNear(result.right.x, expected_right, 1e-5)
def test_bisect_simple(self):
"""Tests that bisect works on a 1 variable scalar valued function."""
wolfe_threshold = 1e-6
x = np.array([0.0, 0.5, 1.0])
y = np.array([1.0, 0.6, 1.2])
dy = np.array([-0.8, 0.6, -0.7])
fun = test_function_x_y_dy(x, y, dy)
val_a = hzl._apply(fun, 0.0) # Value at zero.
val_b = hzl._apply(fun, 1.0) # Value at initial step.
f_lim = val_a.f + (wolfe_threshold * tf.abs(val_a.f))
result = self.evaluate(hzl.bisect(fun, val_a, val_b, f_lim))
self.assertEqual(result.right.x, 0.5)
