def eval_secant2(fun):
val_0 = fun(0.0)
val_a = fun(1.0)
val_b = fun(2.0)
f_lim = val_0.f + (wolfe_threshold * tf.abs(val_0.f))
return self.evaluate(
hzl.secant2(fun, val_0, _interval(val_a, val_b), f_lim))
def _loop_body(curr_interval):
"""The loop body."""
active = ~(curr_interval.converged | curr_interval.failed)
# TODO(b/208441613): Skip updates for batch members that are not active?
secant2_raw_result = hzl.secant2(value_and_gradients_function, val_0,
curr_interval, f_lim,
sufficient_decrease_param,
curvature_param)
secant2_result = HagerZhangLineSearchResult(
## TODO(b/208441613): `& ~curr_interval.failed` should not be needed.
converged=secant2_raw_result.converged & ~curr_interval.failed,
## TODO(b/208441613): `| curr_interval.failed` should not be needed.
failed=secant2_raw_result.failed | curr_interval.failed,
iterations=curr_interval.iterations + tf.cast(active, tf.int32),
func_evals=secant2_raw_result.num_evals,
left=secant2_raw_result.left,
right=secant2_raw_result.right)
should_check_shrinkage = ~(secant2_result.converged
| secant2_result.failed)
def _do_check_shrinkage():
"""Check if interval has shrinked enough."""
old_width = curr_interval.right.x - curr_interval.left.x
new_width = secant2_result.right.x - secant2_result.left.x
sufficient_shrinkage = new_width < old_width * shrinkage_param
func_is_flat = (
_very_close(curr_interval.left.f, curr_interval.right.f)
& _very_close(secant2_result.left.f, secant2_result.right.f))
new_converged = (should_check_shrinkage & sufficient_shrinkage
& func_is_flat)
needs_inner_bisect = should_check_shrinkage & ~sufficient_shrinkage
inner_bisect_args = secant2_result._replace(
converged=secant2_result.converged | new_converged)
def _apply_inner_bisect():
return _line_search_inner_bisection(
value_and_gradients_function, inner_bisect_args,
needs_inner_bisect, f_lim)
return prefer_static.cond(tf.reduce_any(needs_inner_bisect),
_apply_inner_bisect,
lambda: inner_bisect_args)
next_args = prefer_static.cond(tf.reduce_any(should_check_shrinkage),
_do_check_shrinkage,
lambda: secant2_result)
interval_shrunk = (~next_args.failed
& _very_close(next_args.left.x, next_args.right.x))
return [
next_args._replace(converged=next_args.converged | interval_shrunk)
]
def _loop_body(curr_interval):
"""The loop body."""
secant2_raw_result = hzl.secant2(value_and_gradients_function, val_0,
curr_interval, f_lim,
sufficient_decrease_param,
curvature_param)
secant2_result = HagerZhangLineSearchResult(
converged=secant2_raw_result.converged,
failed=secant2_raw_result.failed,
iterations=curr_interval.iterations + 1,
func_evals=secant2_raw_result.num_evals,
left=secant2_raw_result.left,
right=secant2_raw_result.right)
should_check_shrinkage = ~(secant2_result.converged
| secant2_result.failed)
def _do_check_shrinkage():
"""Check if interval has shrinked enough."""
old_width = curr_interval.right.x - curr_interval.left.x
new_width = secant2_result.right.x - secant2_result.left.x
sufficient_shrinkage = new_width < old_width * shrinkage_param
func_is_flat = (
_very_close(curr_interval.left.f, curr_interval.right.f)
& _very_close(secant2_result.left.f, secant2_result.right.f))
new_converged = (should_check_shrinkage & sufficient_shrinkage
& func_is_flat)
needs_inner_bisect = should_check_shrinkage & ~sufficient_shrinkage
inner_bisect_args = secant2_result._replace(
converged=secant2_result.converged | new_converged)
def _apply_inner_bisect():
return _line_search_inner_bisection(
value_and_gradients_function, inner_bisect_args,
needs_inner_bisect, f_lim)
return prefer_static.cond(
tf.reduce_any(input_tensor=needs_inner_bisect),
_apply_inner_bisect, lambda: inner_bisect_args)
next_args = prefer_static.cond(
tf.reduce_any(input_tensor=should_check_shrinkage),
_do_check_shrinkage, lambda: secant2_result)
interval_shrunk = (~next_args.failed
& _very_close(next_args.left.x, next_args.right.x))
return [
next_args._replace(converged=next_args.converged | interval_shrunk)
]
def _loop_body(iteration, found_wolfe, failed, evals, val_left, val_right): # pylint:disable=unused-argument
"""The loop body."""
iteration += 1
secant2_result = hzl.secant2(
value_and_gradients_function,
val_0,
val_left,
val_right,
f_lim,
sufficient_decrease_param=sufficient_decrease_param,
curvature_param=curvature_param)
evals += secant2_result.num_evals
def _failed_fn():
return _LineSearchInnerResult(iteration=iteration,
found_wolfe=False,
failed=True,
num_evals=evals,
left=val_left,
right=val_right)
def _converged_fn():
return _LineSearchInnerResult(iteration=iteration,
found_wolfe=True,
failed=False,
num_evals=evals,
left=secant2_result.left,
right=secant2_result.right)
def _default_fn():
"""Default action."""
new_width = secant2_result.right.x - secant2_result.left.x
old_width = val_right.x - val_left.x
sufficient_shrinkage = new_width < shrinkage_param * old_width
def _sufficient_shrinkage_fn():
"""Action to perform if secant2 shrank the interval sufficiently."""
func_is_flat = (
(tf.math.nextafter(val_left.f, val_right.f) >= val_right.f)
& (tf.math.nextafter(secant2_result.left.f,
secant2_result.right.f) >=
secant2_result.right.f))
is_flat_retval = _LineSearchInnerResult(
iteration=iteration,
found_wolfe=True,
failed=False,
num_evals=evals,
left=secant2_result.left,
right=secant2_result.left)
not_is_flat_retval = _LineSearchInnerResult(
iteration=iteration,
found_wolfe=False,
failed=False,
num_evals=evals,
left=secant2_result.left,
right=secant2_result.right)
return prefer_static.cond(func_is_flat,
true_fn=lambda: is_flat_retval,
false_fn=lambda: not_is_flat_retval)
def _insufficient_shrinkage_fn():
"""Action to perform if secant2 didn't shrink the interval enough."""
update_result = _line_search_inner_bisection(
value_and_gradients_function, secant2_result.left,
secant2_result.right, f_lim)
return _LineSearchInnerResult(iteration=iteration,
found_wolfe=False,
failed=update_result.failed,
num_evals=evals +
update_result.num_evals,
left=update_result.left,
right=update_result.right)
return prefer_static.cond(sufficient_shrinkage,
true_fn=_sufficient_shrinkage_fn,
false_fn=_insufficient_shrinkage_fn)
return prefer_static.case([(secant2_result.failed, _failed_fn),
(secant2_result.converged, _converged_fn)],
default=_default_fn,
exclusive=False)
def eval_secant2(fun, p=lambda x: x):
val_0 = hzl._apply(fun, p(0.0))
val_a = hzl._apply(fun, p(1.0))
val_b = hzl._apply(fun, p(2.0))
f_lim = val_0.f + (wolfe_threshold * tf.abs(val_0.f))
return self.evaluate(hzl.secant2(fun, val_0, val_a, val_b, f_lim))