From 916fcf30c1476e4a4917e9d2f238b30c4594db99 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Fri, 2 Oct 2026 17:31:26 -0500 Subject: [PATCH 01/16] Add a zoom (strong Wolfe) line search --- docs/source/references.bib | 10 + pytensor_ml/optim/line_search.py | 565 +++++++++++++++++++++++++++++++ tests/optim/test_line_search.py | 180 ++++++++++ 3 files changed, 755 insertions(+) create mode 100644 pytensor_ml/optim/line_search.py create mode 100644 tests/optim/test_line_search.py 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/optim/line_search.py b/pytensor_ml/optim/line_search.py new file mode 100644 index 0000000..0fb2f86 --- /dev/null +++ b/pytensor_ml/optim/line_search.py @@ -0,0 +1,565 @@ +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from typing import NamedTuple, Protocol, cast + +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 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 search for a step size along a direction, written as one iteration over a scalar state. + + :func:`search_along` repeats :meth:`step` inside a ``scan`` that stops as soon as it reports it is + done, and every iteration evaluates the loss and its slope at exactly one trial step. + """ + + max_steps: int + + def init(self, value: TensorVariable, slope: TensorVariable) -> State: + """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]: + """Pick one trial step, evaluate it, and return the next state and whether to stop.""" + ... + + def finalize(self, state: State) -> tuple[TensorVariable, TensorVariable]: + """The step size to take and whether the search failed, from the state it stopped in.""" + ... + + +def _cubic_minimizer(a, value_a, slope_a, b, value_b, c, value_c): + """Minimizer of the cubic through three values and the slope at ``a``, or NaN if it has none.""" + 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): + """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 line search for a step satisfying the strong Wolfe conditions, by bracketing and then zooming. + + Built by :func:`zoom_line_search`, which documents the parameters. It follows optax's + ``scale_by_zoom_linesearch``: the bracket grows from the initial guess until it contains a point + satisfying both conditions, then narrows by safeguarded cubic or quadratic interpolation, falling + back to bisection. + """ + + 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) + return { + "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, + } + + def _decrease_error(self, stepsize, value, slope, value0, slope0): + """How far a trial misses sufficient decrease, by Armijo or by the approximate Wolfe test.""" + error = value - value0 - self.slope_rtol * stepsize * slope0 + if self.approx_dec_rtol is not None: + approximate = pt.maximum( + slope - (2 * self.slope_rtol - 1.0) * slope0, + value - value0 - self.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): + """How far a trial misses the strong Wolfe curvature condition.""" + error = pt.maximum(pt.abs(slope) - self.curv_rtol * pt.abs(slope0), 0.0) + return pt.where(pt.isnan(error), np.inf, error) + + def step(self, state, value_and_slope, value0, slope0, guess): + where = pt.where + 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. + bracketing_trial = where(pt.eq(count, 0), guess, self.increase_factor * state["stepsize"]) + if self.max_learning_rate is None: + max_reached = pt.constant(np.array(False)) + else: + max_reached = bracketing_trial >= self.max_learning_rate + bracketing_trial = pt.minimum(bracketing_trial, self.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"], + ) + use_cubic = (cubic > left + 0.2 * width) & (cubic < right - 0.2 * width) + quadratic = _quadratic_minimizer(low, value_low, slope_low, high, value_high) + use_quadratic = ( + ~use_cubic & (quadratic > left + 0.1 * width) & (quadratic < right - 0.1 * width) + ) + middle = where(use_cubic, cubic, where(use_quadratic, quadratic, (low + high) / 2.0)) + zooming = state["interval_found"] + stepsize = where(zooming, middle, bracketing_trial) + + value, slope = value_and_slope(stepsize) + decrease_error = self._decrease_error(stepsize, value, slope, value0, slope0) + curvature_error = self._curvature_error(slope, slope0) + satisfied = pt.maximum(decrease_error, curvature_error) <= self.tol + sufficient_decrease = decrease_error <= self.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"] + too_narrow = (width <= self.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): + # 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 the same + search through optax, which has to be installed for it, 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 Variable + 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." + ) + + 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 = cast( + list[TensorVariable], graph_replace([loss, *gradients], moved, strict=True) + ) + trial_outputs = [trial_value, flat_dot(trial_gradients, directions).astype(dtype)] + 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/optim/test_line_search.py b/tests/optim/test_line_search.py new file mode 100644 index 0000000..f2429e8 --- /dev/null +++ b/tests/optim/test_line_search.py @@ -0,0 +1,180 @@ +from functools import cache + +import numpy as np +import pytensor +import pytensor.tensor as pt +import pytest + +from pytensor.compile.builders import OpFromGraph +from scipy.optimize import line_search as scipy_line_search + +from pytensor_ml.optim.line_search import TRIAL, 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( + "case", + [ + PINNED[name][0] + for name in ("accepts_the_guess", "zooms_into_a_bracket", "grows_until_it_brackets") + ], + ids=["accepts_the_guess", "zooms_into_a_bracket", "grows_until_it_brackets"], +) +def test_the_search_agrees_with_scipys_strong_wolfe_search(case): + """An implementation that shares no code with the port, and lands on the same step wherever the + two searches visit the same trials.""" + 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) + + +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, OpFromGraph) + 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_piece_a_rewrite_removed_is_reported_rather_than_rebuilt(): + """A backend runs the pieces it finds in the rewritten graph and nothing else; when one has gone, it + says so instead of quietly recompiling a copy under some other mode.""" + loss = OBJECTIVES["quadratic"](X) + search = search_along(loss, [X], [pt.grad(loss, X)], [D], zoom_line_search()).step_size.owner.op + + assert find_piece(search, TRIAL).name == TRIAL + with pytest.raises(RuntimeError, match="no 'missing' piece"): + find_piece(search, "missing") From d43559a67a173f80cdcb018ef4df2d4dee8b042e Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Fri, 2 Oct 2026 17:31:38 -0500 Subject: [PATCH 02/16] Run the zoom line search on JAX through optax --- .github/workflows/run_tests.yml | 2 +- pyproject.toml | 1 + pytensor_ml/dispatch/jax/__init__.py | 1 + pytensor_ml/dispatch/jax/line_search.py | 88 +++++++++++++++++++++++++ tests/dispatch/jax/test_line_search.py | 62 +++++++++++++++++ 5 files changed, 153 insertions(+), 1 deletion(-) create mode 100644 pytensor_ml/dispatch/jax/line_search.py create mode 100644 tests/dispatch/jax/test_line_search.py diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml index 115b9b8..21e0b32 100644 --- a/.github/workflows/run_tests.yml +++ b/.github/workflows/run_tests.yml @@ -155,7 +155,7 @@ jobs: - name: Install the backend this job tests run: | conda activate pytensor_ml - if [[ $INSTALL_JAX == "1" ]]; then pip install "jax>=0.8,<0.9.1" jaxlib; fi + if [[ $INSTALL_JAX == "1" ]]; then pip install "jax>=0.8,<0.9.1" jaxlib optax; fi if [[ $INSTALL_MLX == "1" ]]; then pip install "mlx>=0.32.2,<0.33"; fi if [[ $INSTALL_TORCH == "1" ]]; then pip install torch --index-url https://download.pytorch.org/whl/cpu; fi env: diff --git a/pyproject.toml b/pyproject.toml index 7927f1b..967707f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,7 @@ dev = [ "scikit-learn", "mypy", "jax", + "optax", "pyyaml", ] examples = [ 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..7445870 --- /dev/null +++ b/pytensor_ml/dispatch/jax/line_search.py @@ -0,0 +1,88 @@ +import jax +import jax.numpy as jnp + +from pytensor.link.jax.dispatch import jax_funcify + +from pytensor_ml.optim.line_search import TRIAL, LineSearchOp, ZoomLineSearch, find_piece + + +@jax_funcify.register(LineSearchOp) +def jax_funcify_LineSearchOp(op, node=None, **kwargs): + """ + Run the search through optax's, which ``lax.while_loop`` drives, since JAX cannot run a ``scan`` + that stops early. + + optax evaluates the loss through the search's own ``trial`` piece, as the linker rewrote it. It + always starts its search at a step of one, so the direction is scaled by the first trial step on the + way in and the step it finds is scaled back on the way out, with every tolerance measured in steps + scaled to match. + """ + try: + import optax # Optional: only a JAX run of a line search needs it + except ImportError as error: + raise ImportError( + "Running a line search on JAX calls optax's, so optax has to be installed: pip install optax." + ) from error + + search = op.line_search + if not isinstance(search, ZoomLineSearch): + raise NotImplementedError( + f"The JAX backend runs a zoom line search through optax, and has no implementation of " + f"{type(search).__name__}." + ) + kwargs.pop("storage_map", None) + trial = jax_funcify(find_piece(op, TRIAL), **kwargs) + + n_parameters = op.n_parameters + step_dtype, count_dtype = ( + (node.outputs[0].type.dtype, node.outputs[2].type.dtype) + if node is not None + else (None, "int64") + ) + + def line_search(*inputs): + parameters = tuple(inputs[:n_parameters]) + directions = inputs[n_parameters : 2 * n_parameters] + value0, _, guess = inputs[2 * n_parameters : 2 * n_parameters + 3] + others = inputs[2 * n_parameters + 3 :] + one = jnp.ones((), dtype=value0.dtype) + + def value_fn(point): + # The trial at a step of one along the way from the current point to `point` is the loss there + offsets = tuple(there - here for there, here in zip(point, parameters)) + return trial(one, *parameters, *offsets, *others)[0] + + searcher = optax.scale_by_zoom_linesearch( + max_linesearch_steps=search.max_steps, + max_learning_rate=( + None if search.max_learning_rate is None else search.max_learning_rate / guess + ), + tol=search.tol, + increase_factor=search.increase_factor, + slope_rtol=search.slope_rtol, + curv_rtol=search.curv_rtol, + approx_dec_rtol=search.approx_dec_rtol, + stepsize_precision=search.stepsize_precision / guess, + initial_guess_strategy="one", + ) + scaled_directions = tuple(guess * direction for direction in directions) + _, state = searcher.update( + scaled_directions, + searcher.init(parameters), + parameters, + value=value0, + grad=jax.grad(value_fn)(parameters), + value_fn=value_fn, + ) + info = state.info + # optax reports no failure flag. Its errors describe the last trial, which is the step it takes + # whenever it succeeds, so a last trial that misses either condition means the search failed -- + # except when the search stops at the largest step it may take without having bracketed one, + # which it counts as done although the conditions still fail there. + failed = jnp.maximum(info.decrease_error, info.curvature_error) > search.tol + if search.max_learning_rate is not None: + failed = failed & (state.learning_rate < search.max_learning_rate / guess) + step_size = jnp.asarray(state.learning_rate * guess, dtype=step_dtype) + return step_size, failed, jnp.asarray(info.num_linesearch_steps, dtype=count_dtype) + + return line_search diff --git a/tests/dispatch/jax/test_line_search.py b/tests/dispatch/jax/test_line_search.py new file mode 100644 index 0000000..e7de7e3 --- /dev/null +++ b/tests/dispatch/jax/test_line_search.py @@ -0,0 +1,62 @@ +import numpy as np +import pytensor +import pytensor.tensor as pt +import pytest + +pytest.importorskip("jax") +pytest.importorskip("optax") + +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, D, X + + +@pytest.mark.parametrize("case, expected", PINNED.values(), ids=PINNED.keys()) +def test_the_search_runs_through_optax_and_takes_the_same_step(case, expected): + """JAX cannot run the `scan` the search loops in, so it runs optax's search instead. The pinned + values are optax's own, and the default backends' port matches them, so all three agree.""" + 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:] + + +@pytest.mark.parametrize("guess", [0.5, 3.0], ids=["shorter", "longer"]) +def test_a_first_trial_other_than_one_matches_the_default_backend(guess): + """optax always starts at a step of one, so the direction is rescaled by the first trial on the way + in and the step on the way out. A wrong rescaling, of the step or of the tolerances measured in + steps, would land on a different point from the port's.""" + loss = OBJECTIVES["rosenbrock"](X) + [gradient] = pt.grad(loss, [X]) + result = search_along(loss, [X], [gradient], [D], zoom_line_search(), guess=guess) + x0 = np.array([-1.2, 1.0]) + direction = -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 int(on_jax[2]) == int(on_default[2]) + + +def test_a_search_stopped_at_its_largest_step_has_not_failed(): + """Capped at 0.1, the quadratic's search never brackets an acceptable step and stops at the cap, + which counts as done. optax reports no failure flag, and its last trial misses the curvature + condition there, so the flag read off that trial has to make the same exception.""" + loss = OBJECTIVES["quadratic"](X) + [gradient] = pt.grad(loss, [X]) + result = search_along(loss, [X], [gradient], [D], zoom_line_search(max_learning_rate=0.1)) + x0, direction = np.zeros(2), np.ones(2) + + on_jax = pytensor.function([X, D], list(result), mode=jax_mode)(x0, direction) + on_default = pytensor.function([X, D], list(result))(x0, direction) + + assert [np.asarray(value).item() for value in on_jax] == [0.1, False, 1] + assert [np.asarray(value).item() for value in on_default] == [0.1, False, 1] From a916e570862a47e674ca895c5e33252fa8651efe Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Fri, 2 Oct 2026 17:31:45 -0500 Subject: [PATCH 03/16] Run the zoom line search on MLX with a fused Metal kernel --- pytensor_ml/dispatch/mlx/__init__.py | 1 + pytensor_ml/dispatch/mlx/line_search.py | 264 ++++++++++++++++++++++++ tests/dispatch/mlx/test_line_search.py | 71 +++++++ 3 files changed, 336 insertions(+) create mode 100644 pytensor_ml/dispatch/mlx/line_search.py create mode 100644 tests/dispatch/mlx/test_line_search.py 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..3bf7a43 --- /dev/null +++ b/pytensor_ml/dispatch/mlx/line_search.py @@ -0,0 +1,264 @@ +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): + """The fused update for one configuration of the search. Metal compiles it on its 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 is allowed, holding its state once a trial reports that it is done. + + Stopping early would mean reading the stop flag back to Python, which ``mx.compile`` refuses while + it traces, so the loop runs to ``max_steps`` and pays for each trial whether it needs it or not. A + zoom search at single precision on the GPU fuses each trial's bookkeeping into one Metal kernel + between loss evaluations; anything else runs the search's own step graph. + """ + 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" + + 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, kernel): + 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)) + kernel = ( + _zoom_kernel(search) + if isinstance(search, ZoomLineSearch) + and single_precision + and mx.metal.is_available() + and mx.default_device() == mx.gpu + else None + ) + state = stepped(inputs, state) if kernel is None else fused(inputs, state, kernel) + return tuple(finish(*state)) + + return line_search diff --git a/tests/dispatch/mlx/test_line_search.py b/tests/dispatch/mlx/test_line_search.py new file mode 100644 index 0000000..8516d8e --- /dev/null +++ b/tests/dispatch/mlx/test_line_search.py @@ -0,0 +1,71 @@ +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.optim.test_line_search import OBJECTIVES, PINNED + + +@pytest.mark.parametrize("use_compile", [True, False], ids=["compiled", "eager"]) +@pytest.mark.parametrize( + "name", ["accepts_the_guess", "zooms_into_a_bracket", "backs_out_of_a_nan_region"] +) +def test_the_search_takes_the_default_backends_step(name, use_compile): + """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), _ = PINNED[name] + # 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) + # Single precision throughout, which is where the search runs as one Metal kernel per trial; the + # barrier's 0.9 would otherwise promote its graph to float64 + objective_of = { + **OBJECTIVES, + "log_barrier": lambda x: pt.sum((x - np.float32(0.9)) ** 2 - pt.log(1 - x)), + } + loss = objective_of[objective](X) + assert loss.dtype == "float32" + [gradient] = pt.grad(loss, [X]) + result = search_along(loss, [X], [gradient], [D], zoom_line_search(max_steps)) + 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") + + 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")) From 7bbdbbf7bb97947d3a4320f3f601ef7084080887 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Fri, 2 Oct 2026 18:02:51 -0500 Subject: [PATCH 04/16] Replace a typing cast in search_along with a targeted type ignore --- pytensor_ml/optim/line_search.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/pytensor_ml/optim/line_search.py b/pytensor_ml/optim/line_search.py index 0fb2f86..a544e89 100644 --- a/pytensor_ml/optim/line_search.py +++ b/pytensor_ml/optim/line_search.py @@ -1,6 +1,6 @@ from collections.abc import Callable, Sequence from dataclasses import dataclass -from typing import NamedTuple, Protocol, cast +from typing import NamedTuple, Protocol import numpy as np import pytensor @@ -474,7 +474,7 @@ def search_along( 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 Variable + 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. @@ -518,10 +518,11 @@ def search_along( parameter: parameter + trial_step.astype(parameter.dtype) * direction for parameter, direction in zip(parameters, directions) } - trial_value, *trial_gradients = cast( - list[TensorVariable], graph_replace([loss, *gradients], moved, strict=True) - ) - trial_outputs = [trial_value, flat_dot(trial_gradients, directions).astype(dtype)] + 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 From 3c0630b5037d1eaa46126e582b98897a924ae702 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Fri, 2 Oct 2026 18:02:57 -0500 Subject: [PATCH 05/16] Keep a single-precision line search in single precision --- pytensor_ml/optim/line_search.py | 10 ++++++---- tests/optim/test_line_search.py | 21 +++++++++++++++++++++ 2 files changed, 27 insertions(+), 4 deletions(-) diff --git a/pytensor_ml/optim/line_search.py b/pytensor_ml/optim/line_search.py index a544e89..79fb4f1 100644 --- a/pytensor_ml/optim/line_search.py +++ b/pytensor_ml/optim/line_search.py @@ -174,12 +174,14 @@ def step(self, state, value_and_slope, value0, slope0, guess): # 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. - bracketing_trial = where(pt.eq(count, 0), guess, self.increase_factor * state["stepsize"]) + increase_factor = pt.constant(self.increase_factor, dtype=guess.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_reached = bracketing_trial >= self.max_learning_rate - bracketing_trial = pt.minimum(bracketing_trial, self.max_learning_rate) + max_learning_rate = pt.constant(self.max_learning_rate, dtype=guess.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) @@ -200,7 +202,7 @@ def step(self, state, value_and_slope, value0, slope0, guess): ) middle = where(use_cubic, cubic, where(use_quadratic, quadratic, (low + high) / 2.0)) zooming = state["interval_found"] - stepsize = where(zooming, middle, bracketing_trial) + stepsize = where(zooming, middle, bracketing_trial).astype(guess.dtype) value, slope = value_and_slope(stepsize) decrease_error = self._decrease_error(stepsize, value, slope, value0, slope0) diff --git a/tests/optim/test_line_search.py b/tests/optim/test_line_search.py index f2429e8..ba320b8 100644 --- a/tests/optim/test_line_search.py +++ b/tests/optim/test_line_search.py @@ -178,3 +178,24 @@ def test_a_piece_a_rewrite_removed_is_reported_rather_than_rebuilt(): assert find_piece(search, TRIAL).name == TRIAL with pytest.raises(RuntimeError, match="no 'missing' piece"): find_piece(search, "missing") + + +def test_a_single_precision_search_stays_in_single_precision(): + """Hyperparameters float32 cannot hold exactly would otherwise promote the trial step to float64, + which the float32 trial cannot take.""" + 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) + 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) + + assert result.step_size.dtype == "float32" + np.testing.assert_allclose(step_size, expected[0], rtol=1e-5) + assert (bool(failed), int(evaluations)) == (bool(expected[1]), int(expected[2])) From 2003d8abe0ad70660cdc996f25c36198f00ba387 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Fri, 2 Oct 2026 18:03:00 -0500 Subject: [PATCH 06/16] Refuse a line search over a parameter the loss does not read --- pytensor_ml/optim/line_search.py | 11 ++++++++++- tests/optim/test_line_search.py | 8 ++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/pytensor_ml/optim/line_search.py b/pytensor_ml/optim/line_search.py index 79fb4f1..55c481a 100644 --- a/pytensor_ml/optim/line_search.py +++ b/pytensor_ml/optim/line_search.py @@ -9,7 +9,7 @@ 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 truncated_graph_inputs +from pytensor.graph.traversal import ancestors, truncated_graph_inputs from pytensor.scan.utils import until from pytensor.tensor import TensorVariable @@ -501,6 +501,15 @@ def search_along( 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) diff --git a/tests/optim/test_line_search.py b/tests/optim/test_line_search.py index ba320b8..34a79b8 100644 --- a/tests/optim/test_line_search.py +++ b/tests/optim/test_line_search.py @@ -199,3 +199,11 @@ def test_a_single_precision_search_stays_in_single_precision(): assert result.step_size.dtype == "float32" np.testing.assert_allclose(step_size, expected[0], rtol=1e-5) assert (bool(failed), int(evaluations)) == (bool(expected[1]), int(expected[2])) + + +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()) From 1f13493396284e3f3974154a0045002a96c3d112 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Fri, 2 Oct 2026 18:03:13 -0500 Subject: [PATCH 07/16] Run the JAX line search as a while loop over its own step --- .github/workflows/run_tests.yml | 2 +- pyproject.toml | 1 - pytensor_ml/dispatch/jax/line_search.py | 97 +++++++------------------ pytensor_ml/optim/line_search.py | 6 +- tests/dispatch/jax/test_line_search.py | 77 ++++++++++++-------- 5 files changed, 76 insertions(+), 107 deletions(-) diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml index 21e0b32..115b9b8 100644 --- a/.github/workflows/run_tests.yml +++ b/.github/workflows/run_tests.yml @@ -155,7 +155,7 @@ jobs: - name: Install the backend this job tests run: | conda activate pytensor_ml - if [[ $INSTALL_JAX == "1" ]]; then pip install "jax>=0.8,<0.9.1" jaxlib optax; fi + if [[ $INSTALL_JAX == "1" ]]; then pip install "jax>=0.8,<0.9.1" jaxlib; fi if [[ $INSTALL_MLX == "1" ]]; then pip install "mlx>=0.32.2,<0.33"; fi if [[ $INSTALL_TORCH == "1" ]]; then pip install torch --index-url https://download.pytorch.org/whl/cpu; fi env: diff --git a/pyproject.toml b/pyproject.toml index 967707f..7927f1b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,7 +55,6 @@ dev = [ "scikit-learn", "mypy", "jax", - "optax", "pyyaml", ] examples = [ diff --git a/pytensor_ml/dispatch/jax/line_search.py b/pytensor_ml/dispatch/jax/line_search.py index 7445870..c3e46ba 100644 --- a/pytensor_ml/dispatch/jax/line_search.py +++ b/pytensor_ml/dispatch/jax/line_search.py @@ -3,86 +3,39 @@ from pytensor.link.jax.dispatch import jax_funcify -from pytensor_ml.optim.line_search import TRIAL, LineSearchOp, ZoomLineSearch, find_piece +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 through optax's, which ``lax.while_loop`` drives, since JAX cannot run a ``scan`` - that stops early. - - optax evaluates the loss through the search's own ``trial`` piece, as the linker rewrote it. It - always starts its search at a step of one, so the direction is scaled by the first trial step on the - way in and the step it finds is scaled back on the way out, with every tolerance measured in steps - scaled to match. - """ - try: - import optax # Optional: only a JAX run of a line search needs it - except ImportError as error: - raise ImportError( - "Running a line search on JAX calls optax's, so optax has to be installed: pip install optax." - ) from error - - search = op.line_search - if not isinstance(search, ZoomLineSearch): - raise NotImplementedError( - f"The JAX backend runs a zoom line search through optax, and has no implementation of " - f"{type(search).__name__}." - ) + """Run the search's ``step`` in a ``lax.while_loop``, since JAX cannot run a ``scan`` that stops early.""" kwargs.pop("storage_map", None) - trial = jax_funcify(find_piece(op, TRIAL), **kwargs) - - n_parameters = op.n_parameters - step_dtype, count_dtype = ( - (node.outputs[0].type.dtype, node.outputs[2].type.dtype) - if node is not None - else (None, "int64") + 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): - parameters = tuple(inputs[:n_parameters]) - directions = inputs[n_parameters : 2 * n_parameters] - value0, _, guess = inputs[2 * n_parameters : 2 * n_parameters + 3] - others = inputs[2 * n_parameters + 3 :] - one = jnp.ones((), dtype=value0.dtype) - - def value_fn(point): - # The trial at a step of one along the way from the current point to `point` is the loss there - offsets = tuple(there - here for there, here in zip(point, parameters)) - return trial(one, *parameters, *offsets, *others)[0] - - searcher = optax.scale_by_zoom_linesearch( - max_linesearch_steps=search.max_steps, - max_learning_rate=( - None if search.max_learning_rate is None else search.max_learning_rate / guess - ), - tol=search.tol, - increase_factor=search.increase_factor, - slope_rtol=search.slope_rtol, - curv_rtol=search.curv_rtol, - approx_dec_rtol=search.approx_dec_rtol, - stepsize_precision=search.stepsize_precision / guess, - initial_guess_strategy="one", - ) - scaled_directions = tuple(guess * direction for direction in directions) - _, state = searcher.update( - scaled_directions, - searcher.init(parameters), - parameters, - value=value0, - grad=jax.grad(value_fn)(parameters), - value_fn=value_fn, + 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) ) - info = state.info - # optax reports no failure flag. Its errors describe the last trial, which is the step it takes - # whenever it succeeds, so a last trial that misses either condition means the search failed -- - # except when the search stops at the largest step it may take without having bracketed one, - # which it counts as done although the conditions still fail there. - failed = jnp.maximum(info.decrease_error, info.curvature_error) > search.tol - if search.max_learning_rate is not None: - failed = failed & (state.learning_rate < search.max_learning_rate / guess) - step_size = jnp.asarray(state.learning_rate * guess, dtype=step_dtype) - return step_size, failed, jnp.asarray(info.num_linesearch_steps, dtype=count_dtype) + return finish(*final) return line_search diff --git a/pytensor_ml/optim/line_search.py b/pytensor_ml/optim/line_search.py index 55c481a..dba6ecb 100644 --- a/pytensor_ml/optim/line_search.py +++ b/pytensor_ml/optim/line_search.py @@ -467,9 +467,9 @@ def search_along( 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 the same - search through optax, which has to be installed for it, and mlx runs every trial its budget allows, - holding the state once one is accepted. + 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 ---------- diff --git a/tests/dispatch/jax/test_line_search.py b/tests/dispatch/jax/test_line_search.py index e7de7e3..7048083 100644 --- a/tests/dispatch/jax/test_line_search.py +++ b/tests/dispatch/jax/test_line_search.py @@ -4,17 +4,14 @@ import pytest pytest.importorskip("jax") -pytest.importorskip("optax") 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, D, X +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_runs_through_optax_and_takes_the_same_step(case, expected): - """JAX cannot run the `scan` the search loops in, so it runs optax's search instead. The pinned - values are optax's own, and the default backends' port matches them, so all three agree.""" +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]) @@ -28,35 +25,55 @@ def test_the_search_runs_through_optax_and_takes_the_same_step(case, expected): assert (bool(failed), int(evaluations)) == expected[1:] -@pytest.mark.parametrize("guess", [0.5, 3.0], ids=["shorter", "longer"]) -def test_a_first_trial_other_than_one_matches_the_default_backend(guess): - """optax always starts at a step of one, so the direction is rescaled by the first trial on the way - in and the step on the way out. A wrong rescaling, of the step or of the tolerances measured in - steps, would land on a different point from the port's.""" - loss = OBJECTIVES["rosenbrock"](X) - [gradient] = pt.grad(loss, [X]) - result = search_along(loss, [X], [gradient], [D], zoom_line_search(), guess=guess) - x0 = np.array([-1.2, 1.0]) - direction = -pytensor.function([X], gradient)(x0) +# 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, + ), +} - 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 int(on_jax[2]) == int(on_default[2]) - -def test_a_search_stopped_at_its_largest_step_has_not_failed(): - """Capped at 0.1, the quadratic's search never brackets an acceptable step and stops at the cap, - which counts as done. optax reports no failure flag, and its last trial misses the curvature - condition there, so the flag read off that trial has to make the same exception.""" - loss = OBJECTIVES["quadratic"](X) +@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], [D], zoom_line_search(max_learning_rate=0.1)) - x0, direction = np.zeros(2), np.ones(2) + 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) - assert [np.asarray(value).item() for value in on_jax] == [0.1, False, 1] - assert [np.asarray(value).item() for value in on_default] == [0.1, False, 1] + 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])) From 5267ce425c40f71ab7d4b3727e364958bd8a66f9 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Fri, 2 Oct 2026 18:03:20 -0500 Subject: [PATCH 08/16] Return the zoom state in the order the MLX kernel packs it --- pytensor_ml/optim/line_search.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pytensor_ml/optim/line_search.py b/pytensor_ml/optim/line_search.py index dba6ecb..4bdfff8 100644 --- a/pytensor_ml/optim/line_search.py +++ b/pytensor_ml/optim/line_search.py @@ -127,7 +127,7 @@ 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) - return { + state = { "count": pt.constant(0, dtype="int64"), "stepsize": zero, "value": value, @@ -148,6 +148,7 @@ def init(self, value: TensorVariable, slope: TensorVariable) -> State: "safe_stepsize": zero, "safe_value": value, } + return {field: state[field] for field in _ZOOM_FIELDS} def _decrease_error(self, stepsize, value, slope, value0, slope0): """How far a trial misses sufficient decrease, by Armijo or by the approximate Wolfe test.""" From f1b85e5bb62a91bfe66023f170bd592f2d84506d Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Fri, 2 Oct 2026 18:03:39 -0500 Subject: [PATCH 09/16] Choose the MLX line-search path when the function is built --- pytensor_ml/dispatch/mlx/line_search.py | 23 ++++----- tests/dispatch/mlx/test_line_search.py | 66 ++++++++++++++++++++----- 2 files changed, 65 insertions(+), 24 deletions(-) diff --git a/pytensor_ml/dispatch/mlx/line_search.py b/pytensor_ml/dispatch/mlx/line_search.py index 3bf7a43..a9869eb 100644 --- a/pytensor_ml/dispatch/mlx/line_search.py +++ b/pytensor_ml/dispatch/mlx/line_search.py @@ -201,7 +201,8 @@ def mlx_funcify_LineSearchOp(op, node=None, **kwargs): Stopping early would mean reading the stop flag back to Python, which ``mx.compile`` refuses while it traces, so the loop runs to ``max_steps`` and pays for each trial whether it needs it or not. A zoom search at single precision on the GPU fuses each trial's bookkeeping into one Metal kernel - between loss evaluations; anything else runs the search's own step graph. + between loss evaluations; anything else runs the search's own step graph. The choice is made once, + from the default device when the function is built. """ kwargs.pop("storage_map", None) # The pieces as the linker rewrote them, converted by the backend's own OpFromGraph dispatch @@ -218,6 +219,14 @@ def mlx_funcify_LineSearchOp(op, node=None, **kwargs): # 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" + kernel = ( + _zoom_kernel(search) + if isinstance(search, ZoomLineSearch) + and single_precision + and mx.metal.is_available() + and mx.default_device() == mx.gpu + else None + ) def stepped(inputs, state): done = mx.array(False) @@ -227,7 +236,7 @@ def stepped(inputs, state): done = done | stop return state - def fused(inputs, state, kernel): + 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) @@ -250,15 +259,7 @@ def fused(inputs, state, kernel): def line_search(*inputs): value0, slope0 = inputs[2 * n_parameters : 2 * n_parameters + 2] state = list(start(value0, slope0)) - kernel = ( - _zoom_kernel(search) - if isinstance(search, ZoomLineSearch) - and single_precision - and mx.metal.is_available() - and mx.default_device() == mx.gpu - else None - ) - state = stepped(inputs, state) if kernel is None else fused(inputs, state, kernel) + state = stepped(inputs, state) if kernel is None else fused(inputs, state) return tuple(finish(*state)) return line_search diff --git a/tests/dispatch/mlx/test_line_search.py b/tests/dispatch/mlx/test_line_search.py index 8516d8e..c2b1571 100644 --- a/tests/dispatch/mlx/test_line_search.py +++ b/tests/dispatch/mlx/test_line_search.py @@ -15,34 +15,74 @@ 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( - "name", ["accepts_the_guess", "zooms_into_a_bracket", "backs_out_of_a_nan_region"] -) -def test_the_search_takes_the_default_backends_step(name, use_compile): +@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), _ = PINNED[name] + 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) - # Single precision throughout, which is where the search runs as one Metal kernel per trial; the - # barrier's 0.9 would otherwise promote its graph to float64 - objective_of = { - **OBJECTIVES, - "log_barrier": lambda x: pt.sum((x - np.float32(0.9)) ** 2 - pt.log(1 - x)), - } - loss = objective_of[objective](X) + loss = SINGLE_PRECISION[objective](X) assert loss.dtype == "float32" [gradient] = pt.grad(loss, [X]) - result = search_along(loss, [X], [gradient], [D], zoom_line_search(max_steps)) + 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) From fadfcfdbdf241a63f84e919c4438f7eb91827f35 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Fri, 2 Oct 2026 18:03:54 -0500 Subject: [PATCH 10/16] Pin the line search on shared parameters and check it builds a LineSearchOp --- tests/optim/test_line_search.py | 31 +++++++++++++++++++++++++------ 1 file changed, 25 insertions(+), 6 deletions(-) diff --git a/tests/optim/test_line_search.py b/tests/optim/test_line_search.py index 34a79b8..2cdf2bf 100644 --- a/tests/optim/test_line_search.py +++ b/tests/optim/test_line_search.py @@ -5,10 +5,15 @@ import pytensor.tensor as pt import pytest -from pytensor.compile.builders import OpFromGraph from scipy.optimize import line_search as scipy_line_search -from pytensor_ml.optim.line_search import TRIAL, find_piece, search_along, zoom_line_search +from pytensor_ml.optim.line_search import ( + TRIAL, + LineSearchOp, + find_piece, + search_along, + zoom_line_search, +) X = pt.vector("x", dtype="float64") D = pt.vector("d", dtype="float64") @@ -140,7 +145,7 @@ 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, OpFromGraph) + assert isinstance(result.step_size.owner.op, LineSearchOp) assert result.step_size.owner is result.evaluations.owner @@ -169,9 +174,9 @@ def test_hyperparameters_that_cannot_work_are_refused(arguments, message): zoom_line_search(**arguments) -def test_a_piece_a_rewrite_removed_is_reported_rather_than_rebuilt(): - """A backend runs the pieces it finds in the rewritten graph and nothing else; when one has gone, it - says so instead of quietly recompiling a copy under some other mode.""" +def test_a_piece_the_graph_lacks_is_reported_rather_than_rebuilt(): + """A backend runs the pieces it finds in the rewritten graph and nothing else, so asking for one the + graph does not hold raises instead of recompiling a copy under some other mode.""" loss = OBJECTIVES["quadratic"](X) search = search_along(loss, [X], [pt.grad(loss, X)], [D], zoom_line_search()).step_size.owner.op @@ -201,6 +206,20 @@ def test_a_single_precision_search_stays_in_single_precision(): 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) From f3e75623151fd95bb2cb4af46e5e2402adafe1e9 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Mon, 5 Oct 2026 09:09:36 -0500 Subject: [PATCH 11/16] Write the line-search docstrings as one-line imperative summaries --- pytensor_ml/dispatch/jax/line_search.py | 2 +- pytensor_ml/dispatch/mlx/line_search.py | 16 +++++-------- pytensor_ml/optim/line_search.py | 30 ++++++++----------------- 3 files changed, 16 insertions(+), 32 deletions(-) diff --git a/pytensor_ml/dispatch/jax/line_search.py b/pytensor_ml/dispatch/jax/line_search.py index c3e46ba..a984739 100644 --- a/pytensor_ml/dispatch/jax/line_search.py +++ b/pytensor_ml/dispatch/jax/line_search.py @@ -8,7 +8,7 @@ @jax_funcify.register(LineSearchOp) def jax_funcify_LineSearchOp(op, node=None, **kwargs): - """Run the search's ``step`` in a ``lax.while_loop``, since JAX cannot run a ``scan`` that stops early.""" + """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) diff --git a/pytensor_ml/dispatch/mlx/line_search.py b/pytensor_ml/dispatch/mlx/line_search.py index a9869eb..495505b 100644 --- a/pytensor_ml/dispatch/mlx/line_search.py +++ b/pytensor_ml/dispatch/mlx/line_search.py @@ -169,7 +169,7 @@ def _literal(value: float) -> str: @cache def _zoom_kernel(search: ZoomLineSearch): - """The fused update for one configuration of the search. Metal compiles it on its first call.""" + """Build the fused update for one search configuration, which Metal compiles on first call.""" source = _ZOOM_SOURCE.format( n_fields=len(_ZOOM_FIELDS), **_FIELD, @@ -195,15 +195,7 @@ def _zoom_kernel(search: ZoomLineSearch): @mlx_funcify.register(LineSearchOp) def mlx_funcify_LineSearchOp(op, node=None, **kwargs): - """ - Run every trial the search is allowed, holding its state once a trial reports that it is done. - - Stopping early would mean reading the stop flag back to Python, which ``mx.compile`` refuses while - it traces, so the loop runs to ``max_steps`` and pays for each trial whether it needs it or not. A - zoom search at single precision on the GPU fuses each trial's bookkeeping into one Metal kernel - between loss evaluations; anything else runs the search's own step graph. The choice is made once, - from the default device when the function is built. - """ + """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 = ( @@ -219,6 +211,8 @@ def mlx_funcify_LineSearchOp(op, node=None, **kwargs): # 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) @@ -228,6 +222,8 @@ def mlx_funcify_LineSearchOp(op, node=None, **kwargs): 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): diff --git a/pytensor_ml/optim/line_search.py b/pytensor_ml/optim/line_search.py index 4bdfff8..275e3e1 100644 --- a/pytensor_ml/optim/line_search.py +++ b/pytensor_ml/optim/line_search.py @@ -21,17 +21,12 @@ class LineSearch(Protocol): - """ - A search for a step size along a direction, written as one iteration over a scalar state. - - :func:`search_along` repeats :meth:`step` inside a ``scan`` that stops as soon as it reports it is - done, and every iteration evaluates the loss and its slope at exactly one trial step. - """ + """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: - """The state before the first trial, from the loss and its slope at a step of zero.""" + """Return the state before the first trial, from the loss and its slope at a step of zero.""" ... def step( @@ -42,16 +37,16 @@ def step( slope0: TensorVariable, guess: TensorVariable, ) -> tuple[State, TensorVariable]: - """Pick one trial step, evaluate it, and return the next state and whether to stop.""" + """Evaluate one trial step and return the next state and whether to stop.""" ... def finalize(self, state: State) -> tuple[TensorVariable, TensorVariable]: - """The step size to take and whether the search failed, from the state it stopped in.""" + """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): - """Minimizer of the cubic through three values and the slope at ``a``, or NaN if it has none.""" + """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) @@ -63,7 +58,7 @@ def _cubic_minimizer(a, value_a, slope_a, b, value_b, c, value_c): def _quadratic_minimizer(a, value_a, slope_a, b, value_b): - """Minimizer of the quadratic through two values and the slope at ``a``.""" + """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) @@ -94,14 +89,7 @@ def _quadratic_minimizer(a, value_a, slope_a, b, value_b): @dataclass(frozen=True) class ZoomLineSearch: - """ - A line search for a step satisfying the strong Wolfe conditions, by bracketing and then zooming. - - Built by :func:`zoom_line_search`, which documents the parameters. It follows optax's - ``scale_by_zoom_linesearch``: the bracket grows from the initial guess until it contains a point - satisfying both conditions, then narrows by safeguarded cubic or quadratic interpolation, falling - back to bisection. - """ + """A strong Wolfe line search that brackets and then zooms, built by :func:`zoom_line_search`.""" max_steps: int slope_rtol: float @@ -151,7 +139,7 @@ def init(self, value: TensorVariable, slope: TensorVariable) -> State: return {field: state[field] for field in _ZOOM_FIELDS} def _decrease_error(self, stepsize, value, slope, value0, slope0): - """How far a trial misses sufficient decrease, by Armijo or by the approximate Wolfe test.""" + """Return how far a trial misses sufficient decrease, by Armijo or approximate Wolfe.""" error = value - value0 - self.slope_rtol * stepsize * slope0 if self.approx_dec_rtol is not None: approximate = pt.maximum( @@ -163,7 +151,7 @@ def _decrease_error(self, stepsize, value, slope, value0, slope0): return pt.where(pt.isnan(error), np.inf, error) def _curvature_error(self, slope, slope0): - """How far a trial misses the strong Wolfe curvature condition.""" + """Return how far a trial misses the strong Wolfe curvature condition.""" error = pt.maximum(pt.abs(slope) - self.curv_rtol * pt.abs(slope0), 0.0) return pt.where(pt.isnan(error), np.inf, error) From 0e0ababa7b05721c85c374eccaafaeafb1f28ee2 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Mon, 5 Oct 2026 09:09:46 -0500 Subject: [PATCH 12/16] Annotate ZoomLineSearch.step and finalize like the LineSearch protocol --- pytensor_ml/optim/line_search.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/pytensor_ml/optim/line_search.py b/pytensor_ml/optim/line_search.py index 275e3e1..965445a 100644 --- a/pytensor_ml/optim/line_search.py +++ b/pytensor_ml/optim/line_search.py @@ -155,7 +155,14 @@ def _curvature_error(self, slope, slope0): error = pt.maximum(pt.abs(slope) - self.curv_rtol * pt.abs(slope0), 0.0) return pt.where(pt.isnan(error), np.inf, error) - def step(self, state, value_and_slope, value0, slope0, guess): + def step( + self, + state: State, + value_and_slope: ValueAndSlope, + value0: TensorVariable, + slope0: TensorVariable, + guess: TensorVariable, + ) -> tuple[State, TensorVariable]: where = pt.where count = state["count"] low, value_low, slope_low = state["low"], state["value_low"], state["slope_low"] @@ -259,7 +266,7 @@ def step(self, state, value_and_slope, value0, slope0, guess): } return next_state, next_state["done"] | next_state["failed"] - def finalize(self, state): + 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"] & ( From 16377969899a522f78104f59883a402531217d18 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Mon, 5 Oct 2026 09:09:55 -0500 Subject: [PATCH 13/16] Build every zoom search constant in the search's dtype --- pytensor_ml/optim/line_search.py | 34 +++++++++++++++++++++----------- tests/optim/test_line_search.py | 11 +++++++---- 2 files changed, 29 insertions(+), 16 deletions(-) diff --git a/pytensor_ml/optim/line_search.py b/pytensor_ml/optim/line_search.py index 965445a..510efab 100644 --- a/pytensor_ml/optim/line_search.py +++ b/pytensor_ml/optim/line_search.py @@ -140,11 +140,13 @@ def init(self, value: TensorVariable, slope: TensorVariable) -> State: def _decrease_error(self, stepsize, value, slope, value0, slope0): """Return how far a trial misses sufficient decrease, by Armijo or approximate Wolfe.""" - error = value - value0 - self.slope_rtol * stepsize * slope0 + 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 * self.slope_rtol - 1.0) * slope0, - value - value0 - self.approx_dec_rtol * pt.abs(value0), + 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) @@ -152,7 +154,8 @@ def _decrease_error(self, stepsize, value, slope, value0, slope0): def _curvature_error(self, slope, slope0): """Return how far a trial misses the strong Wolfe curvature condition.""" - error = pt.maximum(pt.abs(slope) - self.curv_rtol * pt.abs(slope0), 0.0) + 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( @@ -164,18 +167,19 @@ def step( 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=guess.dtype) + 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=guess.dtype) + 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) @@ -191,20 +195,25 @@ def step( state["cubic_ref"], state["value_cubic_ref"], ) - use_cubic = (cubic > left + 0.2 * width) & (cubic < right - 0.2 * width) + 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 + 0.1 * width) & (quadratic < right - 0.1 * width) + ~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(guess.dtype) + 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) - satisfied = pt.maximum(decrease_error, curvature_error) <= self.tol - sufficient_decrease = decrease_error <= self.tol + 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) @@ -247,7 +256,8 @@ def step( 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"] - too_narrow = (width <= self.stepsize_precision) & (zoom["safe_stepsize"] > 0.0) + 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 = { diff --git a/tests/optim/test_line_search.py b/tests/optim/test_line_search.py index 2cdf2bf..7ddfc1d 100644 --- a/tests/optim/test_line_search.py +++ b/tests/optim/test_line_search.py @@ -8,6 +8,7 @@ from scipy.optimize import line_search as scipy_line_search from pytensor_ml.optim.line_search import ( + STEP, TRIAL, LineSearchOp, find_piece, @@ -186,12 +187,13 @@ def test_a_piece_the_graph_lacks_is_reported_rather_than_rebuilt(): def test_a_single_precision_search_stays_in_single_precision(): - """Hyperparameters float32 cannot hold exactly would otherwise promote the trial step to float64, - which the float32 trial cannot take.""" + """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) + 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) @@ -201,7 +203,8 @@ def test_a_single_precision_search_stays_in_single_precision(): step_size, failed, evaluations = pytensor.function([x, d], list(result))(x0, direction) expected = pytensor.function([X, D], list(in_double))(x0, direction) - assert result.step_size.dtype == "float32" + 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])) From 94581f4cfd1c7301fec4834a8da6948a0c507b72 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Mon, 5 Oct 2026 09:10:04 -0500 Subject: [PATCH 14/16] Pin the default backend's first trial under a guess and a cap --- tests/optim/test_line_search.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/tests/optim/test_line_search.py b/tests/optim/test_line_search.py index 7ddfc1d..869e943 100644 --- a/tests/optim/test_line_search.py +++ b/tests/optim/test_line_search.py @@ -134,6 +134,33 @@ def test_an_accepted_step_satisfies_the_strong_wolfe_conditions(name): 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.""" From 4e0921819c0f5a5bb1d8385c9ad217c31396fb2f Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Mon, 5 Oct 2026 09:10:13 -0500 Subject: [PATCH 15/16] Parametrize the scipy agreement test by case name --- tests/optim/test_line_search.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/tests/optim/test_line_search.py b/tests/optim/test_line_search.py index 869e943..8e09ec1 100644 --- a/tests/optim/test_line_search.py +++ b/tests/optim/test_line_search.py @@ -87,16 +87,12 @@ def test_the_search_takes_the_step_optax_takes(case, expected): @pytest.mark.parametrize( - "case", - [ - PINNED[name][0] - for name in ("accepts_the_guess", "zooms_into_a_bracket", "grows_until_it_brackets") - ], - ids=["accepts_the_guess", "zooms_into_a_bracket", "grows_until_it_brackets"], + "name", ["accepts_the_guess", "zooms_into_a_bracket", "grows_until_it_brackets"] ) -def test_the_search_agrees_with_scipys_strong_wolfe_search(case): +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) From 745ad9d811fb5e30e461e13dfc0eb810e6a22a8f Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Mon, 5 Oct 2026 09:10:29 -0500 Subject: [PATCH 16/16] Drop the find_piece test that the backend tests already cover --- tests/optim/test_line_search.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/tests/optim/test_line_search.py b/tests/optim/test_line_search.py index 8e09ec1..c9da85d 100644 --- a/tests/optim/test_line_search.py +++ b/tests/optim/test_line_search.py @@ -9,7 +9,6 @@ from pytensor_ml.optim.line_search import ( STEP, - TRIAL, LineSearchOp, find_piece, search_along, @@ -198,17 +197,6 @@ def test_hyperparameters_that_cannot_work_are_refused(arguments, message): zoom_line_search(**arguments) -def test_a_piece_the_graph_lacks_is_reported_rather_than_rebuilt(): - """A backend runs the pieces it finds in the rewritten graph and nothing else, so asking for one the - graph does not hold raises instead of recompiling a copy under some other mode.""" - loss = OBJECTIVES["quadratic"](X) - search = search_along(loss, [X], [pt.grad(loss, X)], [D], zoom_line_search()).step_size.owner.op - - assert find_piece(search, TRIAL).name == TRIAL - with pytest.raises(RuntimeError, match="no 'missing' piece"): - find_piece(search, "missing") - - 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,