Add a zoom (strong Wolfe) line search

[jessegrabowski]

Added search_along, which searches for a step size along a direction inside the compiled graph, and zoom_line_search, a strong Wolfe search ported from optax's scale_by_zoom_linesearch. The search is one LineSearchOp node whose inner graph is a scan with until, so it stops at the first trial it accepts. On the default and numba backends it matches optax 0.2.8 exactly on six pinned cases, and agrees with scipy.optimize.line_search where the two visit the same trials.

JAX and MLX cannot run a scan that stops early, so they read the search's named OpFromGraph pieces from the graph their linker has already rewritten. JAX runs the step piece in a lax.while_loop. MLX runs all max_steps trials and holds the state once one is accepted, because mx.compile cannot read the stop flag back while it traces. At float32 on the GPU, MLX runs each trial's bookkeeping as one Metal kernel between loss evaluations.

Nothing calls the search yet. Wiring it into lbfgs is the next step toward #58.

---

📚 Documentation preview 📚: https://pytensor-ml--169.org.readthedocs.build/en/169/