diff --git a/docs/source/references.bib b/docs/source/references.bib index 3ba15b7..1883555 100644 --- a/docs/source/references.bib +++ b/docs/source/references.bib @@ -61,3 +61,13 @@ @article{liu1989lbfgs pages = {503--528}, year = {1989}, } + +@article{hager2005new, + title = {A New Conjugate Gradient Method with Guaranteed Descent and an Efficient Line Search}, + author = {Hager, William W. and Zhang, Hongchao}, + journal = {SIAM Journal on Optimization}, + volume = {16}, + number = {1}, + pages = {170--192}, + year = {2005}, +} diff --git a/pytensor_ml/dispatch/jax/__init__.py b/pytensor_ml/dispatch/jax/__init__.py index b6c1fb5..7046ecd 100644 --- a/pytensor_ml/dispatch/jax/__init__.py +++ b/pytensor_ml/dispatch/jax/__init__.py @@ -2,4 +2,5 @@ # marker op that gets a kernel, mirroring the layout under pytensor_ml/layers. import pytensor_ml.dispatch.jax.attention import pytensor_ml.dispatch.jax.conv +import pytensor_ml.dispatch.jax.line_search import pytensor_ml.dispatch.jax.pooling diff --git a/pytensor_ml/dispatch/jax/line_search.py b/pytensor_ml/dispatch/jax/line_search.py new file mode 100644 index 0000000..a984739 --- /dev/null +++ b/pytensor_ml/dispatch/jax/line_search.py @@ -0,0 +1,41 @@ +import jax +import jax.numpy as jnp + +from pytensor.link.jax.dispatch import jax_funcify + +from pytensor_ml.optim.line_search import FINISH, START, STEP, LineSearchOp, find_piece + + +@jax_funcify.register(LineSearchOp) +def jax_funcify_LineSearchOp(op, node=None, **kwargs): + """Run the search's ``step`` piece in a ``lax.while_loop`` that stops when ``step`` does.""" + kwargs.pop("storage_map", None) + start, step, finish = ( + jax_funcify(find_piece(op, name), **kwargs) for name in (START, STEP, FINISH) + ) + n_parameters = op.n_parameters + max_steps = op.line_search.max_steps + + def line_search(*inputs): + value0, slope0 = inputs[2 * n_parameters : 2 * n_parameters + 2] + initial = tuple(jnp.asarray(field) for field in start(value0, slope0)) + + def keep_going(carry): + iteration, stop, _ = carry + return ~stop & (iteration < max_steps) + + def take_step(carry): + iteration, _, state = carry + *next_state, stop = step(*state, *inputs) + # The loop's carry has to keep its types, which the step's outputs need not match exactly + next_state = tuple( + jnp.asarray(new, dtype=old.dtype) for new, old in zip(next_state, state) + ) + return iteration + 1, jnp.asarray(stop, dtype=bool), next_state + + _, _, final = jax.lax.while_loop( + keep_going, take_step, (jnp.asarray(0), jnp.asarray(False), initial) + ) + return finish(*final) + + return line_search diff --git a/pytensor_ml/dispatch/mlx/__init__.py b/pytensor_ml/dispatch/mlx/__init__.py index 113a295..23f2cc7 100644 --- a/pytensor_ml/dispatch/mlx/__init__.py +++ b/pytensor_ml/dispatch/mlx/__init__.py @@ -3,4 +3,5 @@ import pytensor_ml.dispatch.mlx.attention import pytensor_ml.dispatch.mlx.conv import pytensor_ml.dispatch.mlx.lbfgs +import pytensor_ml.dispatch.mlx.line_search import pytensor_ml.dispatch.mlx.pooling diff --git a/pytensor_ml/dispatch/mlx/line_search.py b/pytensor_ml/dispatch/mlx/line_search.py new file mode 100644 index 0000000..495505b --- /dev/null +++ b/pytensor_ml/dispatch/mlx/line_search.py @@ -0,0 +1,261 @@ +from functools import cache + +import mlx.core as mx + +from pytensor.link.mlx.dispatch import mlx_funcify + +from pytensor_ml.optim.line_search import ( + _ZOOM_FIELDS, + FINISH, + START, + STEP, + TRIAL, + LineSearchOp, + ZoomLineSearch, + find_piece, +) + +_FIELD = {field: index for index, field in enumerate(_ZOOM_FIELDS)} + +# One trial of `ZoomLineSearch.step`, written out for one thread: the state the trial was proposed +# from, then the trial step and the loss and slope there, give the next state and the next trial step. +# It follows the PyTensor graph expression for expression, so the two take the same steps. `maximum` +# and `minimum` propagate NaN, as PyTensor's do and Metal's `fmax`/`fmin` do not. +_ZOOM_HEADER = """ +inline float nan_max(float a, float b) { return (isnan(a) || isnan(b)) ? NAN : fmax(a, b); } +inline float nan_min(float a, float b) { return (isnan(a) || isnan(b)) ? NAN : fmin(a, b); } + +inline float cubic_minimizer(float a, float fa, float fpa, float b, float fb, float c, float fc) { + float db = b - a; + float dc = c - a; + float denominator = (db * dc) * (db * dc) * (db - dc); + float residual_b = fb - fa - fpa * db; + float residual_c = fc - fa - fpa * dc; + float cubic = (dc * dc * residual_b - db * db * residual_c) / denominator; + float quadratic = (-(dc * dc * dc) * residual_b + db * db * db * residual_c) / denominator; + return a + (-quadratic + sqrt(quadratic * quadratic - 3.0f * cubic * fpa)) / (3.0f * cubic); +} + +inline float quadratic_minimizer(float a, float fa, float fpa, float b, float fb) { + float db = b - a; + float curvature = (fb - fa - fpa * db) / (db * db); + return a - fpa / (2.0f * curvature); +} +""" + +_ZOOM_SOURCE = """ + float S[{n_fields}]; + for (uint i = 0; i < {n_fields}; ++i) S[i] = state[i]; + float t = trial[0]; + float f = trial[1]; + float s = trial[2]; + float value0 = anchor[0]; + float slope0 = anchor[1]; + float guess = anchor[2]; + + for (uint i = 0; i < {n_fields}; ++i) next_state[i] = S[i]; + next_trial[0] = t; + if (S[{done}] != 0.0f || S[{failed}] != 0.0f) return; + + float count = S[{count}]; + bool zooming = S[{interval_found}] != 0.0f; + float raw_trial = count == 0.0f ? guess : {increase_factor} * S[{stepsize}]; + bool max_reached = {has_max} ? raw_trial >= {max_learning_rate} : false; + float low = S[{low}], value_low = S[{value_low}], slope_low = S[{slope_low}]; + float high = S[{high}], value_high = S[{value_high}], slope_high = S[{slope_high}]; + float width = fabs(high - low); + + float decrease = f - value0 - {slope_rtol} * t * slope0; + if ({has_approx}) {{ + float approximate = nan_max( + s - (2.0f * {slope_rtol} - 1.0f) * slope0, + f - value0 - {approx_dec_rtol} * fabs(value0) + ); + decrease = nan_min(approximate, decrease); + }} + decrease = nan_max(decrease, 0.0f); + if (isnan(decrease)) decrease = INFINITY; + float curvature = nan_max(fabs(s) - {curv_rtol} * fabs(slope0), 0.0f); + if (isnan(curvature)) curvature = INFINITY; + bool satisfied = nan_max(decrease, curvature) <= {tol}; + bool sufficient = decrease <= {tol}; + bool out_of_steps = count + 1.0f >= {max_steps}; + + float safe_stepsize, safe_value, cubic_ref, value_cubic_ref; + float new_low, new_value_low, new_slope_low, new_high, new_value_high, new_slope_high; + bool interval_found, done, failed; + if (!zooming) {{ + safe_stepsize = sufficient ? t : S[{safe_stepsize}]; + safe_value = sufficient ? f : S[{safe_value}]; + bool high_is_trial = decrease > 0.0f || (f >= S[{value}] && count > 0.0f); + bool low_is_trial = s >= 0.0f && !high_is_trial; + new_low = low_is_trial ? t : S[{stepsize}]; + new_value_low = low_is_trial ? f : S[{value}]; + new_slope_low = low_is_trial ? s : S[{slope}]; + new_high = low_is_trial ? S[{stepsize}] : t; + new_value_high = low_is_trial ? S[{value}] : f; + new_slope_high = low_is_trial ? S[{slope}] : s; + cubic_ref = new_low; + value_cubic_ref = new_value_low; + interval_found = high_is_trial || low_is_trial || satisfied; + done = satisfied || (max_reached && !interval_found); + failed = out_of_steps && !done; + }} else {{ + bool improves = sufficient && f < S[{safe_value}]; + safe_stepsize = improves ? t : S[{safe_stepsize}]; + safe_value = improves ? f : S[{safe_value}]; + done = satisfied; + bool high_is_middle = decrease > 0.0f || f >= value_low; + bool high_is_low = (s * (high - low) >= 0.0f) && !high_is_middle; + float middle_high = high_is_middle ? t : high; + float middle_value_high = high_is_middle ? f : value_high; + float middle_slope_high = high_is_middle ? s : slope_high; + new_high = high_is_low ? low : middle_high; + new_value_high = high_is_low ? value_low : middle_value_high; + new_slope_high = high_is_low ? slope_low : middle_slope_high; + new_low = !high_is_middle ? t : low; + new_value_low = !high_is_middle ? f : value_low; + new_slope_low = !high_is_middle ? s : slope_low; + cubic_ref = (high_is_middle || high_is_low) ? high : low; + value_cubic_ref = (high_is_middle || high_is_low) ? value_high : value_low; + interval_found = true; + bool too_narrow = width <= {stepsize_precision} && safe_stepsize > 0.0f; + failed = (out_of_steps || too_narrow) && !done; + }} + + next_state[{count}] = count + 1.0f; + next_state[{stepsize}] = t; + next_state[{value}] = f; + next_state[{slope}] = s; + next_state[{decrease_error}] = decrease; + next_state[{curvature_error}] = curvature; + next_state[{interval_found}] = interval_found ? 1.0f : 0.0f; + next_state[{done}] = done ? 1.0f : 0.0f; + next_state[{failed}] = failed ? 1.0f : 0.0f; + next_state[{low}] = new_low; + next_state[{value_low}] = new_value_low; + next_state[{slope_low}] = new_slope_low; + next_state[{high}] = new_high; + next_state[{value_high}] = new_value_high; + next_state[{slope_high}] = new_slope_high; + next_state[{cubic_ref}] = cubic_ref; + next_state[{value_cubic_ref}] = value_cubic_ref; + next_state[{safe_stepsize}] = safe_stepsize; + next_state[{safe_value}] = safe_value; + + // The next trial, proposed from the state just written, as the next call of `step` would + float bracketing_trial = {increase_factor} * t; + if ({has_max}) bracketing_trial = fmin(bracketing_trial, {max_learning_rate}); + float next_width = fabs(new_high - new_low); + float left = fmin(new_high, new_low); + float right = fmax(new_high, new_low); + float cubic = cubic_minimizer( + new_low, new_value_low, new_slope_low, new_high, new_value_high, cubic_ref, value_cubic_ref + ); + bool use_cubic = cubic > left + 0.2f * next_width && cubic < right - 0.2f * next_width; + float quadratic = quadratic_minimizer( + new_low, new_value_low, new_slope_low, new_high, new_value_high + ); + bool use_quadratic = + !use_cubic && quadratic > left + 0.1f * next_width && quadratic < right - 0.1f * next_width; + float middle = use_cubic ? cubic : (use_quadratic ? quadratic : (new_low + new_high) / 2.0f); + next_trial[0] = interval_found ? middle : bracketing_trial; +""" + + +def _literal(value: float) -> str: + return "INFINITY" if value == float("inf") else f"{float(value)!r}f" + + +@cache +def _zoom_kernel(search: ZoomLineSearch): + """Build the fused update for one search configuration, which Metal compiles on first call.""" + source = _ZOOM_SOURCE.format( + n_fields=len(_ZOOM_FIELDS), + **_FIELD, + increase_factor=_literal(search.increase_factor), + has_max="true" if search.max_learning_rate is not None else "false", + max_learning_rate=_literal(search.max_learning_rate or 0.0), + slope_rtol=_literal(search.slope_rtol), + curv_rtol=_literal(search.curv_rtol), + has_approx="true" if search.approx_dec_rtol is not None else "false", + approx_dec_rtol=_literal(search.approx_dec_rtol or 0.0), + tol=_literal(search.tol), + stepsize_precision=_literal(search.stepsize_precision), + max_steps=_literal(search.max_steps), + ) + return mx.fast.metal_kernel( + name="pytensor_ml_zoom_line_search", + input_names=["state", "trial", "anchor"], + output_names=["next_state", "next_trial"], + header=_ZOOM_HEADER, + source=source, + ) + + +@mlx_funcify.register(LineSearchOp) +def mlx_funcify_LineSearchOp(op, node=None, **kwargs): + """Run every trial the search allows, holding the state once a trial reports that it is done.""" + kwargs.pop("storage_map", None) + # The pieces as the linker rewrote them, converted by the backend's own OpFromGraph dispatch + start_op, step_op, trial_op, finish_op = ( + find_piece(op, name) for name in (START, STEP, TRIAL, FINISH) + ) + start, step, trial, finish = ( + mlx_funcify(piece, **kwargs) for piece in (start_op, step_op, trial_op, finish_op) + ) + search = op.line_search + n_parameters = op.n_parameters + declared = [variable.type.dtype for variable in finish_op.inner_inputs] + field_dtypes = [getattr(mx, "bool_" if dtype == "bool" else dtype) for dtype in declared] + # Decided from the declared dtype: mlx quietly computes a float64 graph at single precision on the + # GPU, so the arrays reaching the loop would claim float32 for a search declared in float64. + single_precision = declared[_FIELD["value"]] == "float32" + # A zoom search at single precision on the GPU runs each trial's bookkeeping as one Metal kernel + # between loss evaluations, chosen once from the default device when the function is built + kernel = ( + _zoom_kernel(search) + if isinstance(search, ZoomLineSearch) + and single_precision + and mx.metal.is_available() + and mx.default_device() == mx.gpu + else None + ) + + # Stopping early would read the stop flag back to Python, which `mx.compile` refuses while it + # traces, so both loops run all `max_steps` trials + def stepped(inputs, state): + done = mx.array(False) + for _ in range(search.max_steps): + *next_state, stop = step(*state, *inputs) + state = [mx.where(done, held, taken) for held, taken in zip(state, next_state)] + done = done | stop + return state + + def fused(inputs, state): + moved = (*inputs[: 2 * n_parameters], *inputs[2 * n_parameters + 3 :]) + value0, slope0, guess = inputs[2 * n_parameters : 2 * n_parameters + 3] + anchor = mx.stack([value0, slope0, guess]).astype(mx.float32) + packed = mx.stack([field.astype(mx.float32) for field in state]) + stepsize = guess.astype(mx.float32) + if search.max_learning_rate is not None: + stepsize = mx.minimum(stepsize, search.max_learning_rate) + for _ in range(search.max_steps): + value, slope = trial(stepsize, *moved) + packed, next_trial = kernel( + inputs=[packed, mx.stack([stepsize, value, slope]).astype(mx.float32), anchor], + grid=(1, 1, 1), + threadgroup=(1, 1, 1), + output_shapes=[packed.shape, (1,)], + output_dtypes=[mx.float32, mx.float32], + ) + stepsize = next_trial[0] + return [packed[index].astype(dtype) for index, dtype in enumerate(field_dtypes)] + + def line_search(*inputs): + value0, slope0 = inputs[2 * n_parameters : 2 * n_parameters + 2] + state = list(start(value0, slope0)) + state = stepped(inputs, state) if kernel is None else fused(inputs, state) + return tuple(finish(*state)) + + return line_search diff --git a/pytensor_ml/optim/line_search.py b/pytensor_ml/optim/line_search.py new file mode 100644 index 0000000..510efab --- /dev/null +++ b/pytensor_ml/optim/line_search.py @@ -0,0 +1,583 @@ +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from typing import NamedTuple, Protocol + +import numpy as np +import pytensor +import pytensor.tensor as pt + +from pytensor.compile.builders import OpFromGraph +from pytensor.graph.basic import Constant, Variable +from pytensor.graph.replace import graph_replace +from pytensor.graph.traversal import ancestors, truncated_graph_inputs +from pytensor.scan.utils import until +from pytensor.tensor import TensorVariable + +from pytensor_ml.optim.lbfgs import flat_dot +from pytensor_ml.pytensorf.rng import find_generators_drawn_from + +State = dict[str, TensorVariable] +ValueAndSlope = Callable[[TensorVariable], tuple[TensorVariable, TensorVariable]] + + +class LineSearch(Protocol): + """A step-size search written as one trial per :meth:`step`, which :func:`search_along` repeats.""" + + max_steps: int + + def init(self, value: TensorVariable, slope: TensorVariable) -> State: + """Return the state before the first trial, from the loss and its slope at a step of zero.""" + ... + + def step( + self, + state: State, + value_and_slope: ValueAndSlope, + value0: TensorVariable, + slope0: TensorVariable, + guess: TensorVariable, + ) -> tuple[State, TensorVariable]: + """Evaluate one trial step and return the next state and whether to stop.""" + ... + + def finalize(self, state: State) -> tuple[TensorVariable, TensorVariable]: + """Return the step size to take and whether the search failed, from its final state.""" + ... + + +def _cubic_minimizer(a, value_a, slope_a, b, value_b, c, value_c): + """Return the minimizer of the cubic through three values and the slope at ``a``, or NaN.""" + db = b - a + dc = c - a + denominator = (db * dc) ** 2 * (db - dc) + residual_b = value_b - value_a - slope_a * db + residual_c = value_c - value_a - slope_a * dc + cubic = (dc**2 * residual_b - db**2 * residual_c) / denominator + quadratic = (-(dc**3) * residual_b + db**3 * residual_c) / denominator + return a + (-quadratic + pt.sqrt(quadratic**2 - 3.0 * cubic * slope_a)) / (3.0 * cubic) + + +def _quadratic_minimizer(a, value_a, slope_a, b, value_b): + """Return the minimizer of the quadratic through two values and the slope at ``a``.""" + db = b - a + curvature = (value_b - value_a - slope_a * db) / db**2 + return a - slope_a / (2.0 * curvature) + + +_ZOOM_FIELDS = ( + "count", + "stepsize", + "value", + "slope", + "decrease_error", + "curvature_error", + "interval_found", + "done", + "failed", + "low", + "value_low", + "slope_low", + "high", + "value_high", + "slope_high", + "cubic_ref", + "value_cubic_ref", + "safe_stepsize", + "safe_value", +) + + +@dataclass(frozen=True) +class ZoomLineSearch: + """A strong Wolfe line search that brackets and then zooms, built by :func:`zoom_line_search`.""" + + max_steps: int + slope_rtol: float + curv_rtol: float + approx_dec_rtol: float | None + increase_factor: float + max_learning_rate: float | None + tol: float + stepsize_precision: float + + def __post_init__(self) -> None: + if self.max_steps < 1: + raise ValueError(f"max_steps must be at least 1, got {self.max_steps}.") + if not 0.0 < self.slope_rtol < self.curv_rtol < 1.0: + raise ValueError( + f"The Wolfe conditions need 0 < slope_rtol < curv_rtol < 1, got " + f"slope_rtol={self.slope_rtol} and curv_rtol={self.curv_rtol}." + ) + if self.increase_factor <= 1.0: + raise ValueError(f"increase_factor must exceed 1, got {self.increase_factor}.") + + def init(self, value: TensorVariable, slope: TensorVariable) -> State: + zero = pt.zeros((), dtype=value.dtype) + no = pt.constant(np.array(False)) + infinite = pt.constant(np.inf, dtype=value.dtype) + state = { + "count": pt.constant(0, dtype="int64"), + "stepsize": zero, + "value": value, + "slope": slope, + "decrease_error": infinite, + "curvature_error": infinite, + "interval_found": no, + "done": no, + "failed": no, + "low": zero, + "value_low": value, + "slope_low": slope, + "high": zero, + "value_high": value, + "slope_high": slope, + "cubic_ref": zero, + "value_cubic_ref": value, + "safe_stepsize": zero, + "safe_value": value, + } + return {field: state[field] for field in _ZOOM_FIELDS} + + def _decrease_error(self, stepsize, value, slope, value0, slope0): + """Return how far a trial misses sufficient decrease, by Armijo or approximate Wolfe.""" + slope_rtol = pt.constant(self.slope_rtol, dtype=value.dtype) + error = value - value0 - slope_rtol * stepsize * slope0 + if self.approx_dec_rtol is not None: + approx_dec_rtol = pt.constant(self.approx_dec_rtol, dtype=value.dtype) + approximate = pt.maximum( + slope - (2 * slope_rtol - 1.0) * slope0, + value - value0 - approx_dec_rtol * pt.abs(value0), + ) + error = pt.minimum(approximate, error) + error = pt.maximum(error, 0.0) + return pt.where(pt.isnan(error), np.inf, error) + + def _curvature_error(self, slope, slope0): + """Return how far a trial misses the strong Wolfe curvature condition.""" + curv_rtol = pt.constant(self.curv_rtol, dtype=slope.dtype) + error = pt.maximum(pt.abs(slope) - curv_rtol * pt.abs(slope0), 0.0) + return pt.where(pt.isnan(error), np.inf, error) + + def step( + self, + state: State, + value_and_slope: ValueAndSlope, + value0: TensorVariable, + slope0: TensorVariable, + guess: TensorVariable, + ) -> tuple[State, TensorVariable]: + where = pt.where + dtype = guess.dtype + count = state["count"] + low, value_low, slope_low = state["low"], state["value_low"], state["slope_low"] + high, value_high, slope_high = state["high"], state["value_high"], state["slope_high"] + + # Both phases' candidates are cheap scalar arithmetic, so both are formed and the bracket flag + # picks one; only the loss evaluation that follows is expensive. + increase_factor = pt.constant(self.increase_factor, dtype=dtype) + bracketing_trial = where(pt.eq(count, 0), guess, increase_factor * state["stepsize"]) + if self.max_learning_rate is None: + max_reached = pt.constant(np.array(False)) + else: + max_learning_rate = pt.constant(self.max_learning_rate, dtype=dtype) + max_reached = bracketing_trial >= max_learning_rate + bracketing_trial = pt.minimum(bracketing_trial, max_learning_rate) + + width = pt.abs(high - low) + left = pt.minimum(high, low) + right = pt.maximum(high, low) + cubic = _cubic_minimizer( + low, + value_low, + slope_low, + high, + value_high, + state["cubic_ref"], + state["value_cubic_ref"], + ) + cubic_margin = pt.constant(0.2, dtype=dtype) * width + use_cubic = (cubic > left + cubic_margin) & (cubic < right - cubic_margin) + quadratic = _quadratic_minimizer(low, value_low, slope_low, high, value_high) + quadratic_margin = pt.constant(0.1, dtype=dtype) * width + use_quadratic = ( + ~use_cubic + & (quadratic > left + quadratic_margin) + & (quadratic < right - quadratic_margin) + ) + middle = where(use_cubic, cubic, where(use_quadratic, quadratic, (low + high) / 2.0)) + zooming = state["interval_found"] + stepsize = where(zooming, middle, bracketing_trial).astype(dtype) + + value, slope = value_and_slope(stepsize) + decrease_error = self._decrease_error(stepsize, value, slope, value0, slope0) + curvature_error = self._curvature_error(slope, slope0) + tol = pt.constant(self.tol, dtype=dtype) + satisfied = pt.maximum(decrease_error, curvature_error) <= tol + sufficient_decrease = decrease_error <= tol + out_of_steps = (count + 1) >= self.max_steps + trial = (stepsize, value, slope) + + bracketing = {} + bracketing["safe_stepsize"] = where(sufficient_decrease, stepsize, state["safe_stepsize"]) + bracketing["safe_value"] = where(sufficient_decrease, value, state["safe_value"]) + high_is_trial = (decrease_error > 0.0) | ((value >= state["value"]) & (count > 0)) + low_is_trial = (slope >= 0.0) & ~high_is_trial + previous = (state["stepsize"], state["value"], state["slope"]) + bracketing["low"], bracketing["value_low"], bracketing["slope_low"] = ( + where(low_is_trial, new, old) for new, old in zip(trial, previous) + ) + bracketing["high"], bracketing["value_high"], bracketing["slope_high"] = ( + where(low_is_trial, old, new) for new, old in zip(trial, previous) + ) + bracketing["cubic_ref"] = bracketing["low"] + bracketing["value_cubic_ref"] = bracketing["value_low"] + bracketing["interval_found"] = high_is_trial | low_is_trial | satisfied + bracketing["done"] = satisfied | (max_reached & ~bracketing["interval_found"]) + bracketing["failed"] = out_of_steps & ~bracketing["done"] + + zoom = {} + improves_safe = sufficient_decrease & (value < state["safe_value"]) + zoom["safe_stepsize"] = where(improves_safe, stepsize, state["safe_stepsize"]) + zoom["safe_value"] = where(improves_safe, value, state["safe_value"]) + zoom["done"] = satisfied + high_is_middle = (decrease_error > 0.0) | (value >= value_low) + high_is_low = (slope * (high - low) >= 0.0) & ~high_is_middle + high_after_middle = [ + where(high_is_middle, new, old) + for new, old in zip(trial, (high, value_high, slope_high)) + ] + zoom["high"], zoom["value_high"], zoom["slope_high"] = ( + where(high_is_low, new, old) + for new, old in zip((low, value_low, slope_low), high_after_middle) + ) + zoom["low"], zoom["value_low"], zoom["slope_low"] = ( + where(~high_is_middle, new, old) for new, old in zip(trial, (low, value_low, slope_low)) + ) + zoom["cubic_ref"] = where(high_is_middle | high_is_low, high, low) + zoom["value_cubic_ref"] = where(high_is_middle | high_is_low, value_high, value_low) + zoom["interval_found"] = state["interval_found"] + stepsize_precision = pt.constant(self.stepsize_precision, dtype=dtype) + too_narrow = (width <= stepsize_precision) & (zoom["safe_stepsize"] > 0.0) + zoom["failed"] = (out_of_steps | too_narrow) & ~zoom["done"] + + next_state = { + "count": count + 1, + "stepsize": stepsize, + "value": value, + "slope": slope, + "decrease_error": decrease_error, + "curvature_error": curvature_error, + } + for field in bracketing: + next_state[field] = where(zooming, zoom[field], bracketing[field]) + next_state = { + field: pt.as_tensor_variable(value).astype(state[field].dtype) + for field, value in next_state.items() + } + return next_state, next_state["done"] | next_state["failed"] + + def finalize(self, state: State) -> tuple[TensorVariable, TensorVariable]: + # A failed search falls back on the best step that met sufficient decrease, or on not moving at + # all when the last trial was not even finite. + fall_back = state["failed"] & ( + (state["safe_stepsize"] > 0.0) | pt.isinf(state["decrease_error"]) + ) + return pt.where(fall_back, state["safe_stepsize"], state["stepsize"]), state["failed"] + + +def zoom_line_search( + max_steps: int = 20, + *, + slope_rtol: float = 1e-4, + curv_rtol: float = 0.9, + approx_dec_rtol: float | None = 1e-6, + increase_factor: float = 2.0, + max_learning_rate: float | None = None, + tol: float = 0.0, + stepsize_precision: float = 1e-5, +) -> ZoomLineSearch: + r""" + A line search for a step satisfying the strong Wolfe conditions. + + A trial step :math:`t` along a direction :math:`d` from :math:`p` is accepted when it decreases the + loss enough and flattens its slope enough, + + .. math:: + + f(p + t d) \le f(p) + c_1 t \nabla f(p)^\top d, \qquad + |\nabla f(p + t d)^\top d| \le c_2 |\nabla f(p)^\top d|, + + or when it passes the approximate Wolfe test of :cite:t:`hager2005new` in place of the first. The + search stops at the first accepted trial, so it costs one loss and gradient evaluation per trial and + a single one when the initial guess is accepted. When ``max_steps`` trials pass without one, it + takes the best step that met sufficient decrease, or no step at all if the last trial was not + finite. + + Parameters + ---------- + max_steps : int + The most trial steps per search. Default 20. + slope_rtol : float + :math:`c_1`, the fraction of the initial slope the loss has to fall by. Default 1e-4. + curv_rtol : float + :math:`c_2`, the fraction of the initial slope's magnitude the trial slope has to fall below. + Default 0.9. + approx_dec_rtol : float, optional + Relative tolerance of the approximate Wolfe test, or None to require the Armijo condition + alone. Default 1e-6. + increase_factor : float + Growth of the trial step while the search has not yet bracketed an acceptable one. Default 2.0. + max_learning_rate : float, optional + The largest trial step. Default None, which leaves the step unbounded. + tol : float + How far a trial may miss either condition and still be accepted. Default 0.0. + stepsize_precision : float + The bracket width below which a search that has a safe step stops. Default 1e-5. + + Returns + ------- + line_search : ZoomLineSearch + The configured search, for :func:`search_along` or an optimizer's ``line_search`` argument. + """ + return ZoomLineSearch( + max_steps=max_steps, + slope_rtol=slope_rtol, + curv_rtol=curv_rtol, + approx_dec_rtol=approx_dec_rtol, + increase_factor=increase_factor, + max_learning_rate=max_learning_rate, + tol=tol, + stepsize_precision=stepsize_precision, + ) + + +class LineSearchOp(OpFromGraph): + """ + One line search as a single node, built by :func:`search_along`. + + Inputs are the parameters, the directions, the loss and slope at the current point, the first trial + step, and everything else the loss reads; outputs are the step size, whether the search failed, and + how many trials it took. The inner graph runs ``finish(scan(step, start(...)))``, the ``scan`` + stopping at the first trial the search accepts, and every ``step`` evaluates the loss through one + ``trial`` node. Each of those pieces is a named ``OpFromGraph``, so a backend that cannot run the + ``scan`` finds them with :func:`find_piece` in the graph its linker has already rewritten. + + Parameters + ---------- + inputs : list of Variable + The inputs described above, in that order. + outputs : list of Variable + The step size, the failure flag and the trial count. + line_search : LineSearch + The search the inner graph runs. + n_parameters : int + How many parameters, and so how many directions, lead the inputs. + """ + + def __init__( + self, + inputs: list[Variable], + outputs: list[Variable], + *, + line_search: LineSearch, + n_parameters: int, + **kwargs, + ): + super().__init__(inputs, outputs, name="line_search", **kwargs) + self.line_search = line_search + self.n_parameters = n_parameters + + def __eq__(self, other): + return ( + super().__eq__(other) + and self.line_search == other.line_search + and self.n_parameters == other.n_parameters + ) + + def __hash__(self): + return hash((super().__hash__(), self.line_search, self.n_parameters)) + + +# The pieces of a `LineSearchOp`'s inner graph, by name. `start` takes the loss and slope at the current +# point to the first state; `step` takes a state and then every input of the search to the next state and +# whether to stop; `trial` takes a step size, the parameters, the directions and every input after the +# first trial step to the loss and its slope there; `finish` takes the final state to the outputs. +START = "line_search_start" +STEP = "line_search_step" +TRIAL = "line_search_trial" +FINISH = "line_search_finish" + + +def find_piece(op: LineSearchOp, name: str) -> OpFromGraph: + """ + Return the piece of ``op``'s inner graph called ``name``, as the linker left it. + + The pieces are looked for in the inner graph itself and in every inner graph nested in it, so a + piece inside the ``scan`` or inside ``step`` is found too. + + Parameters + ---------- + op : LineSearchOp + A search whose inner graph has been rewritten for the backend converting it. + name : str + One of ``START``, ``STEP``, ``TRIAL`` and ``FINISH``. + + Returns + ------- + piece : OpFromGraph + The piece, with the inner graph the linker rewrote. + """ + pending = [op.fgraph] + while pending: + for node in pending.pop().apply_nodes: + if isinstance(node.op, OpFromGraph) and node.op.name == name: + return node.op + inner = getattr(node.op, "fgraph", None) + if inner is not None: + pending.append(inner) + raise RuntimeError( + f"The line search's inner graph has no {name!r} piece once rewritten for this backend; a " + f"rewrite inlined or merged it, and the backend cannot run the search without it." + ) + + +class SearchResult(NamedTuple): + """ + What a line search found. + + Attributes + ---------- + step_size : TensorVariable + The step to take along the direction. + failed : TensorVariable + Whether the search ran out of trials, or narrowed below its precision, without accepting one. + evaluations : TensorVariable + How many trial steps it evaluated. + """ + + step_size: TensorVariable + failed: TensorVariable + evaluations: TensorVariable + + +def search_along( + loss: TensorVariable, + parameters: Sequence[TensorVariable], + gradients: Sequence[TensorVariable], + directions: Sequence[TensorVariable], + line_search: LineSearch, + guess: float | TensorVariable = 1.0, +) -> SearchResult: + """ + Search for a step size along ``directions`` from the current ``parameters``. + + Each trial evaluates ``loss`` and its gradients with every parameter moved to ``p + t * d``, inside a + ``scan`` that stops at the first trial the search accepts. The loop is one ``OpFromGraph`` node in + the graph it returns. numba and the default backends run that ``scan`` directly, JAX runs its step + in a ``lax.while_loop``, and mlx runs every trial its budget allows, holding the state once one is + accepted. + + Parameters + ---------- + loss : TensorVariable + Scalar loss of the parameters. It has to be deterministic: a loss that draws random numbers would + draw again at every trial, and the trials would not be measuring one function. + parameters : sequence of TensorVariable + The variables the trials move. + gradients : sequence of TensorVariable + The gradient of ``loss`` with respect to each parameter, at the current point. + directions : sequence of TensorVariable + The direction to move each parameter in, one per parameter. + line_search : LineSearch + The search to run, for example :func:`zoom_line_search`. + guess : float or TensorVariable + The first trial step. Default 1.0. + + Returns + ------- + result : SearchResult + The step size, whether the search failed, and how many trials it took. + """ + generators = find_generators_drawn_from([loss]) + if generators: + names = ", ".join(str(generator) for generator in generators) + raise ValueError( + f"A line search evaluates the loss at several trial points, so a loss that draws random " + f"numbers would draw again at each one. This loss draws from {names}; remove the randomness " + f"(dropout, sampling) from the graph the optimizer sees." + ) + + loss_inputs = set(ancestors([loss])) + unused = [parameter for parameter in parameters if parameter not in loss_inputs] + if unused: + names = ", ".join(str(parameter) for parameter in unused) + raise ValueError( + f"The loss does not depend on {names}, so a line search has nothing to measure moving it " + f"by. Leave it out of the parameters searched over." + ) + + dtype = loss.dtype + value0 = loss + slope0 = flat_dot(gradients, directions).astype(dtype) + first_trial = pt.as_tensor_variable(guess).astype(dtype) + + # Each piece is built on stand-ins for what it receives, so it is a graph of its own inputs rather + # than of the point the search starts from. + start_value, start_slope, first_guess = value0.type(), slope0.type(), first_trial.type() + initial = line_search.init(start_value, start_slope) + fields = list(initial) + start = OpFromGraph( + [start_value, start_slope], [initial[field] for field in fields], name=START + ) + + trial_step = first_trial.type() + moved: dict[Variable, Variable] = { + parameter: parameter + trial_step.astype(parameter.dtype) * direction + for parameter, direction in zip(parameters, directions) + } + trial_value, *trial_gradients = graph_replace([loss, *gradients], moved, strict=True) + trial_outputs = [ + trial_value, + flat_dot(trial_gradients, directions).astype(dtype), # type: ignore[arg-type] + ] + moving = [trial_step, *parameters, *directions] + others = [ + variable + for variable in truncated_graph_inputs(trial_outputs, ancestors_to_include=moving) + if variable not in moving and not isinstance(variable, Constant) + ] + trial = OpFromGraph([*moving, *others], trial_outputs, name=TRIAL) + + def value_and_slope(stepsize): + return tuple(trial(stepsize, *parameters, *directions, *others, return_list=True)) + + state = {field: initial[field].type() for field in fields} + next_state, stop = line_search.step( + state, value_and_slope, start_value, start_slope, first_guess + ) + inputs = [*parameters, *directions, value0, slope0, first_trial, *others] + step = OpFromGraph( + [*state.values(), *parameters, *directions, start_value, start_slope, first_guess, *others], + [*(next_state[field] for field in fields), stop], + name=STEP, + ) + + final = {field: initial[field].type() for field in fields} + step_size, failed = line_search.finalize(final) + finish = OpFromGraph(list(final.values()), [step_size, failed, final["count"]], name=FINISH) + + def iterate(*carry_and_inputs): + *next_carry, stop_here = step(*carry_and_inputs, return_list=True) + return next_carry, until(stop_here) + + trace = pytensor.scan( + iterate, + outputs_info=start(value0, slope0, return_list=True), + non_sequences=inputs, + n_steps=line_search.max_steps, + return_updates=False, + ) + outputs = finish.make_node(*(values[-1] for values in trace)).outputs + + search = LineSearchOp(inputs, outputs, line_search=line_search, n_parameters=len(parameters)) + return SearchResult(*search.make_node(*inputs).outputs) diff --git a/tests/dispatch/jax/test_line_search.py b/tests/dispatch/jax/test_line_search.py new file mode 100644 index 0000000..7048083 --- /dev/null +++ b/tests/dispatch/jax/test_line_search.py @@ -0,0 +1,79 @@ +import numpy as np +import pytensor +import pytensor.tensor as pt +import pytest + +pytest.importorskip("jax") + +from pytensor_ml.optim.line_search import search_along, zoom_line_search +from tests.dispatch.jax.test_basic import jax_mode +from tests.optim.test_line_search import OBJECTIVES, PINNED, ROSENBROCK_START, D, X + + +@pytest.mark.parametrize("case, expected", PINNED.values(), ids=PINNED.keys()) +def test_the_search_takes_the_pinned_step_on_jax(case, expected): + objective, x0, direction_of, max_steps = case + loss = OBJECTIVES[objective](X) + [gradient] = pt.grad(loss, [X]) + result = search_along(loss, [X], [gradient], [D], zoom_line_search(max_steps)) + run = pytensor.function([X, D], list(result), mode=jax_mode) + direction = direction_of(pytensor.function([X], gradient)(x0)) + + step_size, failed, evaluations = run(x0, direction) + + np.testing.assert_allclose(float(step_size), expected[0], rtol=1e-9) + assert (bool(failed), int(evaluations)) == expected[1:] + + +# Each row is (objective, start, direction from the gradient, search arguments, first trial, and the +# factor the gradients handed to the search are scaled by). +EDGES = { + "stops_at_the_cap": ( + "quadratic", + np.zeros(2), + lambda g: np.ones(2), + {"max_learning_rate": 0.1}, + 1.0, + 1.0, + ), + "fails_back_onto_the_cap": ( + "quadratic", + np.zeros(2), + lambda g: np.ones(2), + {"max_steps": 1, "max_learning_rate": 5.9}, + 5.9, + 1.0, + ), + "uses_the_gradients_it_is_given": ("rosenbrock", ROSENBROCK_START, lambda g: -g, {}, 1.0, 0.5), + "starts_from_a_nan": ("log_barrier", np.array([1.5]), lambda g: np.ones(1), {}, 1.0, 1.0), + "starts_from_a_trial_other_than_one": ( + "rosenbrock", + ROSENBROCK_START, + lambda g: -g, + {}, + 3.0, + 1.0, + ), +} + + +@pytest.mark.parametrize( + "objective, x0, direction_of, arguments, guess, gradient_scale", + EDGES.values(), + ids=EDGES.keys(), +) +def test_the_search_on_jax_matches_the_default_backend( + objective, x0, direction_of, arguments, guess, gradient_scale +): + loss = OBJECTIVES[objective](X) + [gradient] = pt.grad(loss, [X]) + result = search_along( + loss, [X], [gradient_scale * gradient], [D], zoom_line_search(**arguments), guess=guess + ) + direction = direction_of(pytensor.function([X], gradient)(x0)) + + on_jax = pytensor.function([X, D], list(result), mode=jax_mode)(x0, direction) + on_default = pytensor.function([X, D], list(result))(x0, direction) + + np.testing.assert_allclose(float(on_jax[0]), float(on_default[0]), rtol=1e-9) + assert (bool(on_jax[1]), int(on_jax[2])) == (bool(on_default[1]), int(on_default[2])) diff --git a/tests/dispatch/mlx/test_line_search.py b/tests/dispatch/mlx/test_line_search.py new file mode 100644 index 0000000..c2b1571 --- /dev/null +++ b/tests/dispatch/mlx/test_line_search.py @@ -0,0 +1,111 @@ +import numpy as np +import pytensor +import pytensor.tensor as pt +import pytest + +pytest.importorskip("mlx.core") + +from pytensor.compile.mode import Mode +from pytensor.link.mlx.linker import MLXLinker + +from pytensor_ml.optim.line_search import ( + STEP, + LineSearchOp, + find_piece, + search_along, + zoom_line_search, +) +from tests.dispatch.mlx.test_basic import mx +from tests.optim.test_line_search import OBJECTIVES, PINNED + +# Single precision throughout, which is where the search can run as one Metal kernel per trial; the +# barrier's 0.9 would otherwise promote its graph to float64 +SINGLE_PRECISION = { + **OBJECTIVES, + "log_barrier": lambda x: pt.sum((x - np.float32(0.9)) ** 2 - pt.log(1 - x)), +} + +# The pinned cases, plus the options the pinned cases leave at their defaults, each as (case, search +# arguments, first trial) +CASES = { + **{name: (case, {}, 1.0) for name, (case, _) in PINNED.items()}, + "capped": (PINNED["grows_until_it_brackets"][0], {"max_learning_rate": 0.1}, 1.0), + "armijo_only": (PINNED["zooms_into_a_bracket"][0], {"approx_dec_rtol": None}, 1.0), + "first_trial_not_one": (PINNED["zooms_into_a_bracket"][0], {}, 3.0), +} + +DEVICES = [ + pytest.param( + mx.gpu, + id="kernel", + marks=pytest.mark.skipif(not mx.metal.is_available(), reason="needs Metal"), + ), + pytest.param(mx.cpu, id="step_graph"), +] + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("use_compile", [True, False], ids=["compiled", "eager"]) +@pytest.mark.parametrize("case, arguments, guess", CASES.values(), ids=CASES.keys()) +def test_the_search_takes_the_default_backends_step(case, arguments, guess, use_compile, device): + """mlx runs every trial the search allows and holds the state once one is accepted, so it lands on + the same step, after the same number of trials, as the `scan` that stops there.""" + objective, x0, direction_of, max_steps = case + # Static shapes, as parameters have: a dynamic one adds broadcast checks mlx would drop with a warning + X = pt.vector("x", dtype="float32", shape=x0.shape) + D = pt.vector("d", dtype="float32", shape=x0.shape) + loss = SINGLE_PRECISION[objective](X) + assert loss.dtype == "float32" + [gradient] = pt.grad(loss, [X]) + search = zoom_line_search(max_steps, **arguments) + result = search_along(loss, [X], [gradient], [D], search, guess=guess) + x0 = x0.astype("float32") + direction = direction_of(pytensor.function([X], gradient)(x0)).astype("float32") + mode = Mode(linker=MLXLinker(use_compile=use_compile), optimizer="fast_run") + + with mx.stream(device): + on_mlx = pytensor.function([X, D], list(result), mode=mode)(x0, direction) + on_default = pytensor.function([X, D], list(result))(x0, direction) + + np.testing.assert_allclose(float(on_mlx[0]), float(on_default[0]), rtol=1e-4) + assert (bool(on_mlx[1]), int(on_mlx[2])) == (bool(on_default[1]), int(on_default[2])) + + +def test_a_double_precision_search_takes_the_default_backends_step(): + """Declared in float64, the search keeps to the step graph, which mlx runs at the precision of the + device.""" + (objective, x0, direction_of, max_steps), _ = PINNED["zooms_into_a_bracket"] + X = pt.vector("x", dtype="float64", shape=x0.shape) + D = pt.vector("d", dtype="float64", shape=x0.shape) + loss = OBJECTIVES[objective](X) + [gradient] = pt.grad(loss, [X]) + result = search_along(loss, [X], [gradient], [D], zoom_line_search(max_steps)) + direction = direction_of(pytensor.function([X], gradient)(x0)) + mode = Mode(linker=MLXLinker(), optimizer="fast_run") + + on_mlx = pytensor.function([X, D], list(result), mode=mode)(x0, direction) + on_default = pytensor.function([X, D], list(result))(x0, direction) + + np.testing.assert_allclose(float(on_mlx[0]), float(on_default[0]), rtol=1e-4) + assert (bool(on_mlx[1]), int(on_mlx[2])) == (bool(on_default[1]), int(on_default[2])) + + +def test_the_dispatch_runs_the_pieces_the_compile_mode_rewrote(): + """The pieces are found in the graph the linker rewrote, so they follow the mode the function was + compiled with, rather than being rewritten again under a default of the dispatch's choosing.""" + X = pt.vector("x", dtype="float32", shape=(2,)) + D = pt.vector("d", dtype="float32", shape=(2,)) + loss = pt.sum((X - 3.0) ** 2) / 2 + result = search_along(loss, [X], [pt.grad(loss, X)], [D], zoom_line_search()) + mode = Mode(linker=MLXLinker(), optimizer="fast_run") + + def step_nodes_under(mode): + compiled = pytensor.function([X, D], list(result), mode=mode) + [search] = [ + node.op + for node in compiled.maker.fgraph.apply_nodes + if isinstance(node.op, LineSearchOp) + ] + return len(find_piece(search, STEP).fgraph.apply_nodes) + + assert step_nodes_under(mode) < step_nodes_under(mode.excluding("fusion")) diff --git a/tests/optim/test_line_search.py b/tests/optim/test_line_search.py new file mode 100644 index 0000000..c9da85d --- /dev/null +++ b/tests/optim/test_line_search.py @@ -0,0 +1,242 @@ +from functools import cache + +import numpy as np +import pytensor +import pytensor.tensor as pt +import pytest + +from scipy.optimize import line_search as scipy_line_search + +from pytensor_ml.optim.line_search import ( + STEP, + LineSearchOp, + find_piece, + search_along, + zoom_line_search, +) + +X = pt.vector("x", dtype="float64") +D = pt.vector("d", dtype="float64") + + +def rosenbrock(x, xp=pt): + return xp.sum(100.0 * (x[1:] - x[:-1] ** 2) ** 2 + (1 - x[:-1]) ** 2) + + +OBJECTIVES = { + "quadratic": lambda x, xp=pt: xp.sum((x - 3.0) ** 2) / 2, + "rosenbrock": rosenbrock, + "log_barrier": lambda x, xp=pt: xp.sum((x - 0.9) ** 2 - xp.log(1 - x)), +} + + +@cache +def compiled(objective, max_steps=20): + """One search function per objective and budget, with the point and the direction as inputs, so + every case on an objective runs on one compiled graph.""" + loss = OBJECTIVES[objective](X) + [gradient] = pt.grad(loss, [X]) + result = search_along(loss, [X], [gradient], [D], zoom_line_search(max_steps)) + return pytensor.function([X, D], list(result)), pytensor.function([X], gradient) + + +def search(objective, x0, direction_of, max_steps=20): + run, gradient = compiled(objective, max_steps) + direction = direction_of(gradient(x0)) + step_size, failed, evaluations = run(x0, direction) + return float(step_size), bool(failed), int(evaluations), direction + + +ROSENBROCK_START = np.array([-1.2, 1.0]) + +# Each row is (objective, start, direction from the gradient, budget) and optax 0.2.8's answer from +# `scale_by_zoom_linesearch(max_linesearch_steps=budget, initial_guess_strategy="one")`: the step it +# takes, whether it failed, and how many trials it evaluated. +PINNED = { + "accepts_the_guess": (("quadratic", np.zeros(2), lambda g: np.ones(2), 20), (1.0, False, 1)), + "zooms_into_a_bracket": ( + ("rosenbrock", ROSENBROCK_START, lambda g: -g, 20), + (0.0009383102759, False, 10), + ), + "grows_until_it_brackets": ( + ("rosenbrock", ROSENBROCK_START, lambda g: -g / np.linalg.norm(g) * 1e-3, 20), + (16.0, False, 5), + ), + "backs_out_of_a_nan_region": ( + ("log_barrier", np.zeros(1), lambda g: np.array([2.0]), 20), + (0.1128293312, False, 4), + ), + "fails_uphill_and_keeps_the_last_trial": ( + ("rosenbrock", ROSENBROCK_START, lambda g: g, 20), + (4.031038437e-07, True, 20), + ), + "fails_when_out_of_steps": (("rosenbrock", ROSENBROCK_START, lambda g: -g, 2), (0.5, True, 2)), +} + + +@pytest.mark.parametrize("case, expected", PINNED.values(), ids=PINNED.keys()) +def test_the_search_takes_the_step_optax_takes(case, expected): + """The port follows optax's search trial for trial, so the step, the failure flag and the number of + trials all agree, in the cases that exercise each branch: the guess accepted at once, the bracket + grown, the bracket zoomed, a NaN trial backed out of, and both ways of failing.""" + step_size, failed, evaluations, _ = search(*case) + + np.testing.assert_allclose(step_size, expected[0], rtol=1e-9) + assert (failed, evaluations) == expected[1:] + + +@pytest.mark.parametrize( + "name", ["accepts_the_guess", "zooms_into_a_bracket", "grows_until_it_brackets"] +) +def test_the_search_agrees_with_scipys_strong_wolfe_search(name): + """An implementation that shares no code with the port, and lands on the same step wherever the + two searches visit the same trials.""" + case, _ = PINNED[name] + objective, x0, _, max_steps = case + step_size, _, _, direction = search(*case) + + f = OBJECTIVES[objective] + gradient = pytensor.function([X], pt.grad(OBJECTIVES[objective](X), X)) + expected = scipy_line_search( + lambda x: float(f(x, np)), gradient, x0, direction, c1=1e-4, c2=0.9, maxiter=max_steps + )[0] + + np.testing.assert_allclose(step_size, expected, rtol=1e-9) + + +@pytest.mark.parametrize( + "name", ["zooms_into_a_bracket", "grows_until_it_brackets", "backs_out_of_a_nan_region"] +) +def test_an_accepted_step_satisfies_the_strong_wolfe_conditions(name): + """Checked directly on the objective rather than against another search: sufficient decrease, by + Armijo or by the approximate Wolfe test, and a slope flattened to within ``curv_rtol``.""" + case, _ = PINNED[name] + objective, x0, _, _ = case + step_size, failed, _, direction = search(*case) + f = OBJECTIVES[objective] + gradient = pytensor.function([X], pt.grad(OBJECTIVES[objective](X), X)) + + value0, slope0 = f(x0, np), gradient(x0) @ direction + value, slope = ( + f(x0 + step_size * direction, np), + gradient(x0 + step_size * direction) @ direction, + ) + armijo = value <= value0 + 1e-4 * step_size * slope0 + approximate_wolfe = slope <= (2e-4 - 1) * slope0 and value <= value0 + 1e-6 * abs(value0) + + assert not failed + assert armijo or approximate_wolfe + assert abs(slope) <= 0.9 * abs(slope0) + + +# Each row is (search arguments, first trial) and the step, failure flag and trial count on the quadratic +# from zero along ones, where the slope at a step t is 2 * (t - 3) against -6 at the start. +FIRST_TRIALS = { + "stops_at_the_cap": (({"max_learning_rate": 0.1}, 1.0), (0.1, False, 1)), + "falls_back_onto_the_cap": (({"max_steps": 1, "max_learning_rate": 5.9}, 5.9), (5.9, True, 1)), + "accepts_a_guess_other_than_one": (({}, 3.0), (3.0, False, 1)), +} + + +@pytest.mark.parametrize("case, expected", FIRST_TRIALS.values(), ids=FIRST_TRIALS.keys()) +def test_the_first_trial_follows_the_guess_and_the_cap(case, expected): + """At 0.1 the cap stops a search that has not bracketed, which counts as done. At 5.9 the one trial + allowed misses the curvature condition, so the search fails and falls back on that trial, which met + sufficient decrease. A guess of 3, the minimizer, is accepted where a guess of one would be too.""" + arguments, guess = case + loss = OBJECTIVES["quadratic"](X) + result = search_along( + loss, [X], [pt.grad(loss, X)], [D], zoom_line_search(**arguments), guess=guess + ) + + step_size, failed, evaluations = pytensor.function([X, D], list(result))( + np.zeros(2), np.ones(2) + ) + + assert (float(step_size), bool(failed), int(evaluations)) == expected + + +def test_a_search_from_a_nan_point_takes_no_step(): + """No trial can decrease a loss that is already NaN, and moving anyway would carry the NaN into + the parameters, so the search falls back on a step of zero.""" + step_size, failed, evaluations, _ = search("log_barrier", np.array([1.5]), lambda g: np.ones(1)) + + assert (step_size, failed, evaluations) == (0.0, True, 20) + + +def test_the_search_is_one_node(): + loss = OBJECTIVES["quadratic"](X) + result = search_along(loss, [X], [pt.grad(loss, X)], [D], zoom_line_search()) + + assert isinstance(result.step_size.owner.op, LineSearchOp) + assert result.step_size.owner is result.evaluations.owner + + +def test_a_loss_that_draws_random_numbers_is_refused(): + """Every trial would draw again, so the search would compare values of different functions.""" + generator = pytensor.shared(np.random.default_rng(0), name="noise_rng") + _, noise = pt.random.normal(rng=generator, return_next_rng=True) + loss = pt.sum(X**2) + noise + + with pytest.raises(ValueError, match="draws from noise_rng"): + search_along(loss, [X], [2 * X], [D], zoom_line_search()) + + +@pytest.mark.parametrize( + "arguments, message", + [ + ({"max_steps": 0}, "max_steps must be at least 1"), + ({"slope_rtol": 0.9, "curv_rtol": 0.5}, "0 < slope_rtol < curv_rtol < 1"), + ({"curv_rtol": 1.0}, "0 < slope_rtol < curv_rtol < 1"), + ({"increase_factor": 1.0}, "increase_factor must exceed 1"), + ], + ids=["no_steps", "conditions_swapped", "curvature_unbounded", "no_growth"], +) +def test_hyperparameters_that_cannot_work_are_refused(arguments, message): + with pytest.raises(ValueError, match=message): + zoom_line_search(**arguments) + + +def test_a_single_precision_search_stays_in_single_precision(): + """Hyperparameters and safeguards float32 cannot hold exactly would otherwise promote the search's + arithmetic to float64, deciding its conditions at a precision the float32 backends do not share, + and the float64 trial step that results cannot be fed to the float32 trial at all.""" + x = pt.vector("x", dtype="float32", shape=(2,)) + d = pt.vector("d", dtype="float32", shape=(2,)) + loss = OBJECTIVES["rosenbrock"](x) + search = zoom_line_search(max_learning_rate=0.1, increase_factor=1.7, tol=1e-7) + result = search_along(loss, [x], [pt.grad(loss, x)], [d], search) + loss_in_double = OBJECTIVES["rosenbrock"](X) + in_double = search_along(loss_in_double, [X], [pt.grad(loss_in_double, X)], [D], search) + x0 = ROSENBROCK_START.astype("float32") + direction = np.full(2, 1e-3, dtype="float32") + + step_size, failed, evaluations = pytensor.function([x, d], list(result))(x0, direction) + expected = pytensor.function([X, D], list(in_double))(x0, direction) + + step = find_piece(result.step_size.owner.op, STEP) + assert not [variable for variable in step.fgraph.variables if variable.type.dtype == "float64"] + np.testing.assert_allclose(step_size, expected[0], rtol=1e-5) + assert (bool(failed), int(evaluations)) == (bool(expected[1]), int(expected[2])) + + +def test_the_search_moves_shared_parameters_along_a_direction_of_them(): + """The shape an optimizer builds: shared parameters, and directions that are expressions of them.""" + point = pytensor.shared(ROSENBROCK_START.copy(), name="point") + loss = OBJECTIVES["rosenbrock"](point) + [gradient] = pt.grad(loss, [point]) + result = search_along(loss, [point], [gradient], [-gradient], zoom_line_search()) + + step_size, failed, evaluations = pytensor.function([], list(result))() + + expected = PINNED["zooms_into_a_bracket"][1] + np.testing.assert_allclose(step_size, expected[0], rtol=1e-9) + assert (bool(failed), int(evaluations)) == expected[1:] + + +def test_a_parameter_the_loss_does_not_read_is_refused(): + unused = pt.vector("unused", dtype="float64") + loss = OBJECTIVES["quadratic"](X) + + with pytest.raises(ValueError, match="does not depend on unused"): + search_along(loss, [X, unused], [X - 3.0, unused], [D, unused], zoom_line_search())