Branch-Extended Automatic Differentiation

Introduction

Automatic differentiation (AD) applies the chain rule of calculus to functions expressed as computer programs and is a standard way to obtain derivatives without manually writing or maintaining separate derivative code [1]. For a smooth program, differentiating its execution trace produces the derivative of the represented function (with some caveats in the presence of roundoff, iterative solvers, and other implementation details, see [2] for an overview). Scientific codes, however, often involve continuous but nonsmooth operations such as an absolute value, a pointwise minimum or a pointwise maximum. When such operations appear as primitives in a function, the function typically becomes piecewise smooth; derivatives do not exist in the mathematical sense at inputs that lie exactly on a kink in the piecewise smooth structure. In a computer program, such nonexistence is often represented at a low abstraction level as a branch that switches control flow at such inputs. At higher abstraction levels, this nonexistence occurs at branch points, an intermediate (set of) value(s) at which a primitive enters one of several mathematically valid branches in defining the primal function.

As a motivating example of branches and branch points, consider an implementation of the maximum of two scalars $a$ and $b$:

def max(a, b):
    if(a>=b) return a
    else return b

This code consistently returns $a$ when $a=b$, an arbitrary branch: an equally valid implementation returns $b$ instead:

def max(a, b):
    if(a<=b) return b
    else return a

For every instance where the max primitive appears in a program, the forward execution consistently follows one of two branches. Conventional AD evaluates the derivative of the primal function determined by the corresponding execution trace. However, users of the derivative program, especially practitioners of optimization, may want information about alternative branches not taken by the implementation. We will clarify our intentions later, but note that the execution trace associated with alternative branches may have vastly different derivatives [3, 4]. A user may also want derivatives corresponding to alternative branches that were nearly followed within a prescribed tolerance; that is, for small $\mathrm{eps}>0$, replacing the first implementation of the max primitive with

def max(a, b):
    if(a + eps >=b) return a
    else return b

would have admitted an alternative branch for the same function input.

To overcome this limitation of conventional AD, we introduce branch-extended automatic differentiation (BEAD). At a given function input, BEAD records the branch taken at supported nonsmooth primitives such as abs and max, identifies alternative branches within a prescribed tolerance, and replays the program once per path through branches—an assignment of one branch at every branch point—thereby locking those branches. Each replay defines a smooth branch extension at the query input, differentiable by ordinary JAX AD [5]. The output is a set of branch-extension gradients rather than only the gradient of the default path. BEAD returns tolerance- and program-dependent branch-extension gradients; these need not form a classical generalized derivative, as discussed in § 2.2.

Our implementation and test cases are mostly motivated by manifold sampling, a derivative-free optimization method that relies on composite objective functions of a certain form, discussed further in § 2. However, BEAD’s applicability is not specific to composite objectives or derivative-free optimization, and we demonstrate some preliminary results with scalar-valued nonsmooth functions in the context of the constrained nonsmooth solver PyGRANSO [6, 7] in § 4.2. BEAD supports both local tolerances on intermediate values and input-scaled tolerances based on first-order input sensitivities, as described in § 3.2.

The contributions of this work are a record-and-replay formulation of BEAD, a JAX implementation for array operations that commonly appear in definitions of piecewise smooth functions, implementations of two branch-discovery rules, and an evaluation against manual implementations of branch-extension gradients from past work.

Background

We now review piecewise-smooth programs and existing generalized-derivative concepts. We then review BEAD’s connection to nonsmooth optimization methods.

Piecewise-smooth programs

A continuous piecewise-$C^1$ function is locally represented by finitely many smooth selection functions, with one or more selection functions active at a switching surface [8]. For input queries off the switching surfaces, ordinary AD of the execution trace returns the derivative of the locally unique selection function. For input queries on the switching surface, a single execution trace corresponds to only one selection function; it does not describe the other selection functions that determine nearby first-order behaviors of the piecewise-smooth function.

In related work, abs-normal form makes this switching structure explicit by representing piecewise-smooth functions using smooth operations connected by absolute-value switches [9]. It supports systematic piecewise-linearization and generalized algorithmic differentiation. BEAD is closely related but does not transform a function into abs-normal form, nor does it return a complete piecewise-linear local model of the function. Instead, it records the branch taken at each supported JAX primitive, identifies within-tolerance alternative branches, extends them to the current input, and applies ordinary AD to each, producing branch-extension values and gradients.

We list several other related works in piecewise-smooth programs here. Intensional derivatives formalize the derivatives computed by AD for piecewise-analytic programs and agree with ordinary derivatives almost everywhere [3]. Conservative fields provide a chain-rule-compatible set-valued calculus for broad classes of nonsmooth programs and support convergence analyses for stochastic-gradient methods [4, 10]. Automatic subdifferentiation seeks to compute a mathematically valid generalized subderivative for a qualified program [11]. These approaches primarily justify and select generalized derivative information, but do not provide information about alternative branches, which is required for robust nonsmooth optimization solvers that motivate this work.

Relationship to generalized derivatives

Several notions of generalized derivative exist [12, 13]. These generalized derivatives include Dini or Hadamard directional derivatives that describe a one-sided rate of change in a specified direction. The Bouligand, or B-, subdifferential $\partial_B f(x)$ collects limits of Jacobians from differentiable points in a neighborhood of a given point. For locally Lipschitz scalar functions, the Clarke subdifferential $\partial_C f(x)$ and Clarke generalized directional derivative represent the convex hull of limiting gradients and their support function, respectively. Convex subgradients are defined by global supporting inequalities for the epigraph of a convex function.

In general, BEAD computes none of these generalized derivatives. Given an input query $x$, BEAD produces a finite set

where $\tau> 0$ is an absolute tolerance used in determining alternate branches, $\mathcal{P}_{\tau}(x)$ is the set of recorded paths through branches, and $f_p$ is the branch extension obtained by locking path $p$. The set $\mathcal{P}_{\tau}(x)$ depends on the program, supported primitives, and tolerance strategy. § 3.3 gives theoretical conditions under which $\mathcal{G}_0(x)$ contains $\partial_B f(x)$.

Moreover, with a positive tolerance $\tau> 0$, inactive smooth extensions are deliberately included in $\mathcal{P}_{\tau}(x)$, and so it is generally not the case when $\tau> 0$ that $\mathcal{G}_{\tau}(x)$ coincides with any of these generalized derivatives.

Another related approach is lexicographic differentiation [14], which computes directional derivatives for a provided sequence of infinitesimal perturbations. Lexicographic differentiation has been employed to find solutions of both systems of nonsmooth equations and parametric differential equations [15]. The directional nature of lexicographic differentiation distinguishes it from BEAD, which instead discovers a set of paths according to a finite tolerance and differentiates each corresponding branch extension separately.

Relationship to nonsmooth optimization methods

Various methods have been proposed for the solution of nonsmooth optimization problems, that is, optimization problems for which the objective function and possibly some constraints are nonsmooth functions of their inputs. The most robust of these methods generally employ a generalized derivative at points of nondifferentiability, or else construct approximations of generalized derivatives from conventional gradients at slightly perturbed points in neighborhoods of points of nondifferentiability. Bundle methods [16, 17, 18] use subgradients at various query points in the optimization domain to build and update cutting-plane models of objectives and constraints. Gradient sampling methods [19, 20] evaluate conventional gradients at randomly sampled points of differentiability near a given input and computes the gradients’ convex hull to approximate a Clarke subdifferential.

In contrast to the mechanisms of bundle methods and gradient sampling methods, both of which aggregate (sub)gradient information across multiple input points, BEAD considers a single input point and differentiates branch extensions identified from program structure. This mode of computation is particularly useful within the paradigm of manifold sampling [21, 22, 23, 24], which seeks to identify (a subset of) active manifolds defining a continuous selection at any given incumbent. Our motivating manifold-sampling problems have the composite form

where $F:\mathbb{R}^{n}\to\mathbb{R}^{p}$ is expensive to evaluate and $h:\mathbb{R}^{p}\to\mathbb{R}$ is a known, possibly nonsmooth scalar mapping. Manifold sampling builds separate models of the components of $F$ and uses values and gradients associated with active manifolds of $h$; “activity” here generally corresponds to branch extensions. BEAD automates the path-identification and differentiation interface required by manifold sampling methods. We discuss benchmark problems further in § 4.

However, BEAD can be applied beyond manifold sampling methods. As a particular example, a finite BEAD set $\mathcal{G}_{\mathcal{T}}(x)$ can be employed to construct a cutting-plane model that would appear in a bundle method. The downstream bundle method may minimize the support function $\max_{g\in\mathcal{G}_{\mathcal{T}}(x)}g^{T}d$ when choosing a direction for use in a line-search or in a trust-region subproblem. This use case does not require a decomposition $f=h\circ F$ like manifold sampling is designed to exploit; instead, BEAD can be applied directly to a piecewise-smooth function and generate BEAD sets of practical utility.

Implementation

We now BEAD’s record-and-replay execution model, its branch-discovery rules, batched and compiled replay, and its user-facing API.

Design and execution model

BEAD is implemented as a Python package built on JAX [5]. BEAD does not replace or modify JAX’s differentiation machinery in any significant way. Instead, BEAD changes the behavior of the forward pass, generating multiple branch extensions that can then be differentiated normally. User code defining the primal function must import a NumPy-like wrapper module for a set of supported nonsmooth operations, as well as jax.numpy in order to define smooth operations.

A recording pass executes the user’s primal function once. At each supported nonsmooth primitive, BEAD records the executed branch and identifies alternative branches that are within a specified tolerance, to be defined in § 3.2. The implementation of BEAD distinguishes several types of primitives. Reduction extrema primitives, such as min and max, produce a scalar output and record the selected index together with other indices within tolerance of the extreme. A unary elementwise primitive, such as abs, records a sign branch at entries within tolerance of zero. Binary elementwise extrema primitives, such as minimum and maximum, record an operand branch at each array index. These branches are stored in small typed records and can have large downstream effects: elementwise primitives can induce exponentially many branch combinations when many entries are within tolerance simultaneously, whereas reduction extrema primitives produce a number of branches that grows, at most, linearly in the size of the input vector. These combinatorial issues should be taken into account when designing primal functions.

The recorded branches, that is the executed branches and alternative branches within tolerance taken together, define a PathSet. A single resolved path through branches within a PathSet fixes one branch at every branch point in the program; downstream code can refer to it by a path label string. The number of paths in a PathSet is generally exponential in the number of branch points that admit within-tolerance alternatives for a given input. To avoid instantiating an exponential number of paths in memory whenever possible, our implementation uniquely identifies each path by a single integer, which represents the position of that path in a sequence defined by the Cartesian product of the branches available at each branch point encountered in a program. A PathSet by default therefore only needs to store the available branches at each branch point separately, as opposed to storing all possible combinations.

PathSet objects also support set operations such as union and intersection. Such operations may destroy the Cartesian-product structure, in which case BEAD stores an explicit list of paths. Whether stored compactly or explicitly, evaluating every path requires work proportional to the number of paths, which is exponential in the number of branch points that admit alternatives within tolerance. Thus, users of BEAD may need to control the runtime or memory use of this approach either by judiciously reducing the tolerance or designing their problem in a way that ensures that the number of alternatives is sufficiently bounded. Future work might investigate, for example, automatic sampling strategies in BEAD in order to limit resource use when user-controlled remedies are apparently insufficient.

A replay pass executes the user program using the branches fixed by one resolved path in a PathSet. Thus, a replay pass changes the forward values produced by the supported nonsmooth primitives. Once those branches are locked, the replay pass is an ordinary smooth JAX program for that branch extension, wherever its elementary operations are smooth. BEAD can consequently apply jax.grad or jax.value_and_grad directly to the replayed function. In particular, BEAD does not implement a separate reverse pass or a custom derivative rule for each composite nonsmooth function; by changing the forward branches, JAX may compute the derivative of that branch extension automatically.

Figure 1 summarizes this execution model. Branch discovery is performed during forward execution, while each replay is an ordinary JAX computation that can be differentiated and compiled.

Branch discovery and path-set growth

Branch discovery is controlled by absolute and relative tolerances. We propose two explicit rules for branch discovery employing these tolerances. Every supported primitive has an associated scalar switching condition $q$, computed from the intermediate value(s) that primitive takes. The sign of $q$ determines the branch taken at that branch point. For example, abs switches on the sign of its single input, while maximum switches on the sign of the difference between its two operands. BEAD applies each primitive’s $q$ consistently during both recording and replay of branches. Because each branch point is tested independently, a combination of branches recorded in $P_\tau(x)$ need not correspond to any single, jointly reachable input perturbation; the path set could include branches that can’t be reached by any nearby input; Theorem 3.2 is a small instance. The local branch-discovery rule admits the alternative branch when $|q|$ is below a tolerance expressed in the units of that intermediate quantity. We implement such a rule as

where $v$ is an operation-specific scale: the selected extremum for a reduction, the larger-magnitude operand for a binary extremum, and the input to an absolute value.

Such a test is generally inexpensive and easy to implement, but may not always behave in a way that is intuitive to a user; intermediate values and their sensitivity to function input perturbations may have vastly different scales from one another. As an alternative, it may be more intuitive to ask whether perturbing the function input within some small tolerance could flip the sign of $q$. In other words, we search for alternative program paths that would be taken for some input within a small neighborhood around the current function input $x$. Exact propagation of such a neighborhood through a nonlinear program, however, is a challenging problem. Interval methods can propagate bounds through a program, but may overestimate reachable ranges. Taylor models retain higher-order information at greater cost [25, 26].

Noting these challenges, BEAD employs an input-scaled branch-discovery rule to serve as a first-order approximation to the propagation of a neighborhood through a nonlinear program. Because $q$ is computed from the function input $x$ through the forward pass, it can be viewed as a scalar function $q(x)$ of $x$, with switching surface $\{x : q(x) = 0\}$. For an input ball $\lVert\delta\rVert\le\tau$, the first-order range of the switching quantity is controlled by $\tau\lVert\nabla q(x)\rVert_\ast$, where $\lVert\cdot\rVert_\ast$ is the dual norm. Forward-mode AD propagates input sensitivities through the program; BEAD uses the resulting magnitudes to estimate this scale and adjust the local tolerance. The sensitivities only guide branch discovery; replay gradients are still computed by ordinary JAX AD. This sensitivity-based approximation does not provide any strong bounds, as it deliberately neglects higher-order effects for nonlinear functions. Like the local rule, the input-scaled rule is a practical branch-discovery rule and generally does not correspond precisely to a generalized derivative.

Nonetheless, the first-order input-scaled rule corrects two failure modes of the local rule in the linear examples below. More generally, if $q$ has an $L$-Lipschitz gradient on the input ball, then

The input-scaled rule models the linear term and neglects the remainder. If the switching quantity $q(x)$ contracts input perturbations, the local rule may admit a branch whose switching surface is outside the input ball. If $q(x)$ expands input perturbations, the local rule may reject a branch whose switching surface lies inside the ball. The input-scaled rule corrects both decisions for these linear examples. We illustrate an example of these behaviors in Figure 2.

As an example of the mitigation strategies mentioned above, we present a policy specific to the abs function applied entrywise to a vector. When the vector contains $k$ entries that are near zero within tolerance, there are $2^k$ recorded paths, enumerated by a $k$-fold Cartesian product of sign branches $\pm$. Naively, one could employ a policy, called via abs_policy="enumerate" in BEAD, that causes BEAD to record both signs $\pm$ for each ambiguous entry. However, with the default abs_policy="zero", BEAD instead retains one compact representative branch whose recorded value is simply the absolute value, and the associated derivative is taken at zero at any ambiguous entry within tolerance $\tau$. This zero is not one of the two one-sided derivatives $\pm1$ of abs; rather, it is the minimum-norm element of its Clarke subdifferential $[-1,1]$ at a tie, the standard representative when a single (sub)gradient is required, for example in a stationarity measure. The default policy avoids $2^k$ recorded paths in favor of one path, at the expense of a less informative $\mathcal{G}_{\tau}(x)$.

The absolute-value policy separates three decisions: which alternatives are discovered, how they are represented, and which resolved paths are evaluated. Other primitives with many within-tolerance alternatives could make the same separation; discovery, representation, and evaluation need not all be exhaustive. Future work involves topics such as lazy path generation, sampling, limits on retained paths, removal of duplicate or nearly identical gradients, compact representations of independent or convex derivative sets, and optimizer-guided selection of paths relevant to a search direction. Such policies should be studied in terms of downstream accuracy (e.g., in stationarity tests or optimization search directions), memory use, and runtime. Currently, abs is the only primitive that has such a policy implemented in BEAD, because other primitives did not demonstrate path-set explosion in our tests. Future work will study similar policies, for example, for the maximum and minimum primitive. These primitives in particular exhibit an explosion in the size of discovered path-sets when two compared input vectors have many indices that are equal in value or nearly equal in value.

Generalized derivatives at zero tolerance

We consider the special case of $G_\tau(x)$ where $\tau=0$. Every recorded path in $P_0(x)$ fixes one branch at each branch point, and so defines a smooth branch extension $f_p$ that agrees with $f$ at $x$. These branch extensions are selection functions in the sense of Scholtes [8]. To formalize a result about $G_0(x)$, we call a path (selection function) $p\in P_0(x)$ essentially active if $x$ lies in the closure of the interior of $\{y:f(y)=f_p(y)\}$.

Proof. A recorded path fixes a branch at every branch point, so each $f_p$ is a composition of continuously differentiable operations and is $C^1$ in some open neighborhood of $x$. Let $U_1$ be the intersection of these finitely many neighborhoods.

We will show that $f$ is a selection of the family $\{f_p\}$ near $x$. Consider the execution of the program defining $f$. At any branch point whose switching quantity $q$ is nonzero at the intermediate values entering the branch point (call these values $v(x)$), the same branch is executed for all values in an open neighborhood of $v(x)$. By continuity of $f$, this means the same branch will be executed for all $y$ in some open neighborhood of $x$. Denote the intersection of these finitely many (corresponding to each branch point with nonzero intermediate values) open sets by $U_2$. On the other hand, at any branch point whose switching quantity $q$ vanishes at $v(x)$, every alternative is recorded by assumption. Consider the open set $U := U_1 \cap U_2$. We have both that $x \in U$ and for every $y \in U$, the branches executed by the program at $y$ belong to some path $p \in P_0(x)$, that is, $f(y) = f_p(y)$. Thus $f$ is a continuous selection of the finite family of $C^1$ functions $\{f_p\}_{p\in P_0(x)}$ on $U$.

[8 Prop. 4.3.1] gives that $\partial_B f(x)$ is the set of gradients at $x$ of the essentially active members of this family, and that $\partial_C f(x)$ is its convex hull. Because every path in the full product is evaluated, $G_0(x) = \{\nabla f_p(x) : p \in P_0(x)\}$, which contains the gradients of the essentially active paths. This is the first inclusion in (2); taking convex hulls gives the second. Finally, $G_0(x) = \partial_B f(x)$ exactly when every $\nabla f_p(x)$ equals the gradient of some essentially active path, which is the stated criterion. Equality of the sets implies equality of their convex hulls. □

The following simple one-dimensional example shows the hypotheses of Theorem 3.1 can hold while equality in eq. (2) fails.

The hypotheses of Theorem 3.1 hold in BEAD when atol=rtol=0, tol_mode="local", and abs_policy="enumerate". With the default abs_policy="zero", they fail: $G_0(0) = \{0\}$ for $f(t) = |t|$, against $\partial_B f(0) = \{-1, 1\}$. The guarantees of Theorem 3.1 also fail for $\tau> 0$: a recorded branch need no longer agree with any executed branch at $x$, and so $f_p(x) \ne f(x)$ is possible; the extensions then no longer form a selection family for $f$.

Note that BEAD does not test essential activity, which Theorem 3.1 requires. For programs built from smooth operations and absolute values, the linear independence kink qualification (LIKQ) of abs-normal form [27 Def. 2] is a sufficient condition: under it, every sign combination at the active switches is realized on a set of positive measure [28 Cor. 7], so every recorded path is essentially active and equality holds in (2).

Batched and compiled replay

JAX supports JIT (just-in-time) compilation to accelerate the execution speed of functions and their gradient computation. JIT compilation introduces a tradeoff. The first call to JAX incurs tracing and compilation costs, whereas later calls can reuse the compiled program and avoid repeated Python dispatch. JAX usually needs to re-compile programs when their control flow or certain array dimensions change. To avoid frequent recompilations during the interactions of an optimization method with BEAD, our implementation pads the size of some data structures to the next power of two and removes padded results afterward. This performs some unused work but reduces the number of compiled batch shapes. Compiled replay functions are cached by the user function, operation layout, differentiated arguments, and auxiliary-output option. This optimization generally reduces replay overhead but is independent of the possibly combinatorial enumeration of recorded paths.

We call a collection of resolved paths that are evaluated together a path batch; a path batch is evaluated in batched replay. For a collection of paths with the same sequence of branch points, BEAD first validates their structure and converts their corresponding branches into dense JAX arrays. BEAD then applies jax.value_and_grad to one replay, uses jax.vmap to evaluate the resolved paths together, and applies jax.jit to compile the vectorized calculation.

User-facing API

The public interface follows JAX’s functional-transformation style. That is, each BEAD transformation accepts a Python function and returns a callable with the same user arguments. Supported nonsmooth operations are written using branch_extended_AD.numpy; smooth operations may use jax.numpy. Table 1 summarizes the layers of the API.

User goal	Transformation
Discover within-tolerance paths	record
Evaluate one resolved path	replay
Differentiate the default path	grad, value_and_grad
Differentiate one resolved path	replay_grad, replay_value_and_grad
Differentiate all discovered paths	all_value_and_grad
Evaluate known paths in one compiled call	replay_value_and_grad_batch

Listing 1 provides a demonstration of recording and replay. A recording pass returns the value for the default path along with a PathSet. Replay accepts one resolved path and returns the corresponding branch-extension value.

Listing 1: Recording and replaying within-tolerance paths.

import jax.numpy as jnp
import branch_extended_AD as bead
import branch_extended_AD.numpy as bnp

def h(z):
    return bnp.sum(bnp.maximum(z, 0.0))

z = jnp.array([0.0, 1.0, -2.0])
value, paths = bead.record(h, atol=1.0e-8)(z)

for path in paths:
    locked_value = bead.replay(h, path)(z)

Derivative transformations use the same path objects. The default path routines return the default-path derivative together with the paths discovered at the provided input. Replay routines evaluate one known path. When every recorded path is required, all_value_and_grad performs discovery and replay in one call.

Listing 2: Computing branch-extension gradients from resolved paths.

(value, default_grad), paths = bead.value_and_grad(h, atol=1.0e-8)(z)
branch_data = [bead.replay_value_and_grad(h, path)(z) for path in paths]
all_branch_data, paths = bead.all_value_and_grad(h, atol=1.0e-8)(z)

The options atol, rtol, tol_mode, and abs_policy dictate the scope of branch discovery. The options argnums and has_aux follow the corresponding JAX transformations. If a user employing BEAD has already collected resolved paths, they may use replay_value_and_grad_batch to evaluate the same resolved paths through the compiled batched replay rather than needing to rediscover them.

Test cases

We now discuss several use cases designed to probe the capabilities of BEAD, and note the capabilities’ alignment with our design intentions.

Manifold sampling

All manifold-sampling test cases use the composite form (1). A traditional derivative-free method will generally build and maintain one surrogate $m_f$ to model the full objective $f = h \circ F$. Manifold sampling instead builds and maintains separate models $m_1,\ldots,m_p$, aggregated in the vector-valued model $m_F$, for the components $F_1,\ldots,F_p$ of $F$ and combines these $p$ models using the known structure of $h$. The componentwise construction $h \circ m_F$ often provides a more accurate local model of $f$ than $m_f$.

When $h$ is nonsmooth, a manifold sampling method requires both function value and gradients for its active manifolds. For a rigorous definition of active manifold, we refer the reader again to [21, 22, 23, 24]. As a motivating example of the definition for the sake of understanding, for $h(z) = \lVert z\rVert_1$, manifold identification reduces to labeling the signs of the entries of $F(x)$. For a more complicated mapping, such as the censored-$\ell_1$ loss

manual identification and differentiation are substantially more involved. BEAD automates this problem-specific work and reduces the risk of errors.

To test whether our AD-encoded $h$ functions are correct, we needed a test suite broad enough to exercise BEAD’s various tolerances and branch discovery policies. Rather than hand-selecting a handful of problem/$h$ combinations, we ran every smooth vector-valued function $F$ from a benchmark set against every $h$ we had previously hand-coded in the IBCDFO package [29].

For $F$, we use the full set of 53 test problems from the Moré–Wild DFO benchmark collection [30]; problems include a reference starting point $x_0$. These problems span a low range of dimensions $(2 \le n \le12)$ and output dimensions $2 \le p \le45$. The Moré–Wild benchmark is intentionally small in scale, because it is intended to represent the types of problems typically encountered in simulation-based optimization, where $F$ is very expensive to evaluate and so a single simplex gradient approximation (requiring $n+1$ function evaluations) is cost-prohibitive.

As for $h$, to really test the capabilities of BEAD, we defined nine different $h$ functions, listed in Table 2. For many of these $h$ functions (e.g., piecewise min/max functions) optimization regularly encounters points of nondifferentiability, making branch discovery of critical importance. Other $h$ functions (e.g., the piecewise-quadratic and censored-$\ell_1$ losses) are compositions with so many branching locations that hand-differentiation can be impractical and error-prone.

Name	Formula
one_norm	$h(z)=\sum_{i=1}^{p}|z_i|$
pw_maximum	$h(z)=\max\{z_1,\ldots,z_p\}$
pw_maximum_sq	$h(z)=\max\{z_1^2,\ldots,z_p^2\}$
pw_minimum	$h(z)=\min\{z_1,\ldots,z_p\}$
pw_minimum_sq	$h(z)=\min\{z_1^2,\ldots,z_p^2\}$
quantile	$q$-th quantile of $\{z_1^2,\ldots,z_p^2\}$
max_plus_qvp	$\max\{z_1,\ldots,z_{p-1}\}+\alpha\sum_{i\geq p}\max\{z_i,0\}^2$
censored_L1_loss	$h(z;c,d)=\sum_{i=1}^{p}|d_i-\max(z_i,c_i)|$
piecewise_quadratic	$h(z;Q_j,z_j,c_j)=\max_j\left\{\lVert z-z_j\rVert_{Q_j}^2+c_j\right\}$

Two of the $h$ functions in Table 2 deserve discussion. In our tests, h_quantile is hardcoded to $q = 1$, which makes it functionally identical to h_pw_minimum_squared; we therefore exclude h_quantile from the counted instances below to avoid redundancies. We keep h_quantile in Table 2 only to document the general quantile-based $h$ function.

Similarly, h_max_plus_quadratic_violation_penalty has its penalty weight $\alpha$ hardcoded to 0, and so it reduces to h_pw_maximum restricted to the first $p-1$ components; we choose to keep this $h(z)$ in the test set (and in the instance count) because it exercises a different code path even though its values coincide with h_pw_maximum.

The last two entries in Table 2 are not fixed functions but function factories: create_censored_L1_loss_hfun builds a one-sided censored $\ell_1$ loss from problem-specific data $(c,d)$, following Womersley’s censored discrete linear $\ell_1$ approximation setup [31], and create_piecewise_quadratic_hfun builds a max-of-quadratics surrogate from problem-specific $(Q_j,z_j,c_j)$ triples, as used in [23]. For both, we reuse the exact problem data that accompanied those papers’ own benchmark studies (the first random instance/seed for each of the 53 problems), rather than generating our own. This yields $53\times8=424$ benchmark problem instances.

Finally, we comment that many of these $h$ functions are unbounded for many of the $F$ instances on an unconstrained domain. For example, $h(z)=\min\{z_1,\ldots,z_p\}$ composed with an unconstrained $F$ has no finite lower bound. To address this, we impose box constraints $\ell\leq x\leq u$ around $x_0$ for all problems. We employ a simple, deterministic box centered at $x_0$, with a per-coordinate radius of

so $\ell_i=x_{0,i}-r_i$ and $u_i=x_{0,i}+r_i$. The max with 20 matters because several Moré–Wild starting points have components at or near 0, and so a purely multiplicative radius would collapse to (almost) zero; scaling by $|x_{0,i}|$ matters because other starting points have components in the thousands, where a fixed radius would be far too tight to let a solver explore the domain. The function $F$ is not defined for all values in the above box-constrained domain. We observed overflows in $F$ for Meyer’s function (corresponding to problem 17) and Osborne’s second function (problems 36 and 37). For these three problems, we use a tighter box ($r_i=\max(5,|x_{0,i}|)$) to guarantee that $F$ stays finite, instead of employing the radius we prescribed above.

For the IBCDFO test cases, we created an integration module that wraps the BEAD interface and adapts values, gradient matrices, and path labels to IBCDFO’s existing $h$-function interface. The PyGRANSO experiments in the following Section use a similar adapter.

Other nonsmooth optimization methods

To verify that BEAD is useful beyond manifold sampling and the explicit composition $h(F(x))$, we also coupled its branch-extension gradients to PyGRANSO, a BFGS–SQP method for general nonsmooth constrained optimization from the NCVX software package [6, 7]. We replaced the single-gradient term in the local search model by the support function of the BEAD gradient set and used the corresponding worst-case directional derivative in the weak-Wolfe line search. We retained the original BFGS-SQP method as a baseline and also tested an intermediate variant that used BEAD gradients only in the stationarity test.

We used two small problems for which all within-tolerance paths can be enumerated. The first is the constrained nonsmooth Rosenbrock example

initialized at the infeasible kink $(1,1)$. Both the baseline BFGS-SQP method and the stationarity-only variant failed before accepting an iterate. Our extension employing BEAD gradients reached a feasible stationary point in 18 iterations and 42 function evaluations.

As a second example, we consider a three-variable one-dimensional total-variation denoising problem (the fused-lasso signal approximator)

Starting from $(1,1,1)$, the baseline required 15 iterations and 67 evaluations to terminate at a point near the solution $(2/3,2/3,2/3)$, using BEAD only for stationarity required 15 iterations and 43 evaluations, and our variant of the method required one iteration and three evaluations. This behavior can likely be explained by the structure of the problem; the initial point is already on the same manifold as the solution to the optimization problem, on which the selection function corresponding to the recorded path is quadratic. These two problems illustrate how BEAD’s path and branch-extension information ports to a general nonsmooth solver; with only two problems, they are meant to be illustrative and we do not claim they provide a broad performance comparison between the baseline and BEAD-based variants.

Results

We now compare BEAD against hand-coded derivatives to see understand how it affects optimizer performance and then study its runtime.

Hand-coded versus BEAD $h$ functions

As a direct consistency check, the automated gamma-example test compares function values and gradient matrices at generic and tie inputs using relative tolerance $10^{-6}$ and absolute tolerance $10^{-8}$. The full benchmark below provides a broader downstream comparison across all eight pairs of $h$ functions. The remaining question is whether this equivalence translates to downstream performance: does employing BEAD $h$ functions in a manifold sampling algorithm, instead of the hand-coded ones, change the behavior of the optimization method significantly? We ran every $(\text{problem},h)$ pair from the previous subsection twice – once with each version of $h$ – through the same manifold sampling solver, and compared the results using data profiles [30].

A data profile plots, for each solver, the fraction of problems solved as a function of computational budget, measured in simplex gradients (i.e., groups of $n+1$ function evaluations, where $n$ is the number of variables for a given problem). We consider a problem instance to be solved after $k$ evaluations when the objective value $f(x_k)$ first satisfies

where $f(x_0)$ is the starting value, $f_{\mathrm{best}}$ is the smallest objective value found by any solver on that instance across all evaluations, and $\kappa> 0$ is the convergence gate; gate governs the quality of the final optimizer result, while tolerance governs the proximity of input or intermediate values that determines which branches are taken during function evaluation. The gate is a purely function-value-based criterion, measuring how close each solver comes to the best observed value, and can be computed without gradient information.

In these experiments we use the local branch-discovery rule with rtol = 0, so the reported $\tau$ values are the API parameter atol, the absolute tolerance controlling which alternative branches are considered. We investigated three values: $\tau= 10^{-8}$ (our baseline, the tightest), $10^{-6}$ (looser), and $10^{-4}$ (much looser). This parameter directly affects BEAD’s runtime, because a looser tolerance means BEAD must track more within-tolerance paths. The set of discovered paths also affects the downstream optimizer behavior in a problem-dependent and scale-dependent way. A good choice of $\tau$ is therefore specific to $h$ and to how $h$ is implemented; the three values with which we experiment are chosen to illustrate this sensitivity, not to prescribe a general value of $\tau$ useful for all imaginable $h$ functions.

Figure 3 shows data profiles across all three tolerance values and both convergence gates. At the loose gate $\kappa= 10^{-3}$ (panel (d)), the two curves are visually indistinguishable regardless of $\tau$, solving about 93% of the 424 problem instances within the iteration budget; the three tolerance settings produce essentially the same profile, so only one representative is shown. At the tighter gate $\kappa= 10^{-7}$ (panels (a)–(c)), differences become visible: at $\tau= 10^{-8}$, BEAD and hand-coded performance are comparable; as $\tau$ is loosened, the BEAD curve drops below hand-coded, reflecting the looser branch-discovery tolerance. Hand-coded performance is largely insensitive to $\tau$ across this range.

Vector (PDF)

The end result is that replacing a hand-derived $h$ with its BEAD counterpart does not meaningfully degrade the performance of manifold sampling at the loose gate or at our tightest tolerance $\tau= 10^{-8}$; performance degrades mildly at the tightest gate as $\tau$ is loosened, but mirrors the coarser branch-discovery tolerance. This robustness check supports using automatically generated BEAD path and branch-extension gradient information in place of hand-coded derivatives.

Timing

In this section, we evaluate the runtime of BEAD, compared with the hand-coded gradients. BEAD’s branch discovery incurs bookkeeping overhead in each evaluation, and its current implementation records branches on every call rather than compiling it once with JAX’s JIT.

Table 3 reports the median wall-clock time per function evaluation for each $h$, across three values of the branch-discovery tolerance $\tau$. Both hand-coded and BEAD derivatives are run at the same matched $\tau$ in each column group, and all timings are restricted to runs that terminated normally (ran the full evaluation budget) for both versions, so the comparison is over the same amount of work.

$h$ function	$\tau= 10^{-8}$ hand	$\tau= 10^{-8}$ BEAD	$\tau= 10^{-8}$ ratio	$\tau= 10^{-6}$ hand	$\tau= 10^{-6}$ BEAD	$\tau= 10^{-6}$ ratio	$\tau= 10^{-4}$ hand	$\tau= 10^{-4}$ BEAD	$\tau= 10^{-4}$ ratio
censored_L1_loss	9.66	35.74	3.70	7.31	40.86	5.59	14.60	35.97	2.46
piecewise_quadratic	7.94	25.16	3.17	13.16	24.53	1.86	5.84	33.10	5.67
max_plus_qvp	7.74	30.39	3.93	7.81	34.12	4.37	7.17	30.92	4.31
one_norm	11.26	30.07	2.67	13.75	38.22	2.78	15.74	45.24	2.87
pw_maximum	7.02	23.29	3.31	7.72	25.10	3.25	7.08	30.08	4.25
pw_maximum_sq	8.56	24.00	2.80	8.65	29.11	3.36	6.97	36.15	5.18
pw_minimum	4.29	15.29	3.56	6.29	20.88	3.32	3.00	25.56	8.53
pw_minimum_sq	5.07	18.11	3.57	5.55	22.61	4.07	4.05	26.06	6.43
overall (median)	7.84	24.58	3.14	7.76	27.11	3.49	7.03	32.01	4.56

At $\tau= 10^{-8}$, the BEAD derivative is slower than the hand-coded derivative by a median factor of about $3\times$. As $\tau$ is loosened from $10^{-8}$ to $10^{-6}$ to $10^{-4}$, the median ratio increases ($3.1\times\rightarrow3.5\times\rightarrow4.6\times$). Both implementations may return more gradients at looser tolerances, but their work differs: the hand-coded routines construct gradient columns directly, whereas BEAD records additional paths and replays each branch extension through JAX AD. The table shows that hand-coded median times remain comparatively stable while BEAD times increase. In the simulation-based optimization settings this benchmark represents (§ 4), $F$ is far more expensive to evaluate than any computation of derivatives of paths through $h$ and so this overhead is negligible.

Conclusions

We introduced BEAD, a Python package that collects gradient information for a set of paths through programs describing piecewise-differentiable functions evaluated at or near kinks. The package uses JAX, a popular Python framework, for the differentiation of each individual path.

We demonstrate the usefulness of BEAD on two benchmarks: a manifold-sampling benchmark where hand-coded manifold-identification logic was replaced with BEAD, resulting in comparable solver performance for objectives of the form $h(F(x))$; and PyGRANSO examples, showing that the same path and branch-extension gradient information can support a general nonsmooth optimizer.

The set BEAD returns is not in general a classical generalized derivative; Theorem 3.1 gives conditions on the program and the enumeration policy under which it contains the Bouligand subdifferential.

The main limitation of BEAD (as well as its hand-coded alternatives) is path-set growth. If a program contains multiple branch points with nearby alternatives (for example, $\max(a, b)$ with $a \approx b$), the total number of paths grows exponentially with the number of such branch points. We present one strategy specific to the abs function to mitigate this path explosion, but future work is needed to develop similar policies for other primitives, and evaluate them in terms of computational cost and the gradient information that they preserve.

Acknowledgments

This work was supported in part by the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research, Scientific Discovery through Advanced Computing (SciDAC) Program through the FASTMath Institute under Contract No. DE-AC02-06CH11357.

The initial version of the BEAD software prototype was written by hand, but subsequently refactored, and its API and internal data structures redesigned with the help of AI coding agents. The test cases and experiments were selected by the authors, and testing and plotting scripts developed with the help of AI agents. AI agents were also used in the preparation of early drafts of the user manual and the manuscript, and subsequently heavily revised manually by the authors. The authors take full responsibility for the contents of the manuscript and software submission.

Software availability. The BEAD source code, tests, and examples are available in the project repository [32] under the MIT License.