diff --git a/desc/derivatives.py b/desc/derivatives.py
index 5122222efc..49d409a399 100644
--- a/desc/derivatives.py
+++ b/desc/derivatives.py
@@ -26,7 +26,7 @@ class _Derivative(ABC):
     """
 
     @abstractmethod
-    def __init__(self, fun, argnum=0, mode=None, **kwargs):
+    def __init__(self, fun, argnum=0, mode=None, has_aux=False, **kwargs):
         pass
 
     @abstractmethod
@@ -118,6 +118,13 @@ class AutoDiffDerivative(_Derivative):
         ``'hess'`` (Hessian of a scalar function),
         or ``'jvp'`` (Jacobian vector product)
         Default = ``'fwd'``
+    chunk_size : int
+        Will calculate the Jacobian
+        ``chunk_size`` columns at a time, instead of all at once.
+    has_aux : bool
+        Indicates whether fun returns a pair where the first element is considered
+        the output of the mathematical function to be differentiated and the second
+        element is auxiliary data. Default False.
 
     Raises
     ------
@@ -125,11 +132,14 @@ class AutoDiffDerivative(_Derivative):
 
     """
 
-    def __init__(self, fun, argnum=0, mode="fwd", chunk_size=None, **kwargs):
+    def __init__(
+        self, fun, argnum=0, mode="fwd", chunk_size=None, has_aux=False, **kwargs
+    ):
 
         self._fun = fun
         self._argnum = argnum
         self._chunk_size = chunk_size
+        self._has_aux = has_aux
         self._set_mode(mode)
 
     def compute(self, *args, **kwargs):
@@ -152,7 +162,7 @@ def compute(self, *args, **kwargs):
         return self._compute(*args, **kwargs)
 
     @classmethod
-    def compute_vjp(cls, fun, argnum, v, *args, **kwargs):
+    def compute_vjp(cls, fun, argnum, v, *args, has_aux=False, **kwargs):
         """Compute v.T * df/dx.
 
         Parameters
@@ -178,12 +188,15 @@ def compute_vjp(cls, fun, argnum, v, *args, **kwargs):
         _ = kwargs.pop("rel_step", None)  # unused by autodiff
 
         def _fun(*args):
+            if has_aux:
+                f, aux = fun(*args, **kwargs)
+                return v.T @ f, aux
             return v.T @ fun(*args, **kwargs)
 
-        return jax.grad(_fun, argnum)(*args)
+        return jax.grad(_fun, argnum, has_aux=has_aux)(*args)
 
     @classmethod
-    def compute_jvp(cls, fun, argnum, v, *args, **kwargs):
+    def compute_jvp(cls, fun, argnum, v, *args, has_aux=False, **kwargs):
         """Compute df/dx*v.
 
         Parameters
@@ -215,11 +228,16 @@ def _fun(*x):
                 _args[i] = xi
             return fun(*_args, **kwargs)
 
-        y, u = jax.jvp(_fun, tuple(args[i] for i in argnum), v)
-        return u
+        out = jax.jvp(_fun, tuple(args[i] for i in argnum), v, has_aux=has_aux)
+        # out is either y, u or y, u, aux. We only want u and aux if present
+        if len(out) == 2:
+            return out[1]  # don't want to return a tuple if no aux
+        return out[1:]  # slice out y
 
     @classmethod
-    def compute_jvp2(cls, fun, argnum1, argnum2, v1, v2, *args, **kwargs):
+    def compute_jvp2(
+        cls, fun, argnum1, argnum2, v1, v2, *args, has_aux=False, **kwargs
+    ):
         """Compute d^2f/dx^2*v1*v2.
 
         Parameters
@@ -255,14 +273,18 @@ def compute_jvp2(cls, fun, argnum1, argnum2, v1, v2, *args, **kwargs):
             argnum2 = tuple([i + 1 for i in argnum2])
             v2 = tuple(v2)
 
-        dfdx = lambda dx1, *args: cls.compute_jvp(fun, argnum1, dx1, *args, **kwargs)
+        dfdx = lambda dx1, *args: cls.compute_jvp(
+            fun, argnum1, dx1, *args, has_aux=has_aux, **kwargs
+        )
         d2fdx2 = lambda dx1, dx2: cls.compute_jvp(
-            dfdx, argnum2, dx2, dx1, *args, **kwargs
+            dfdx, argnum2, dx2, dx1, *args, has_aux=has_aux, **kwargs
         )
         return d2fdx2(v1, v2)
 
     @classmethod
-    def compute_jvp3(cls, fun, argnum1, argnum2, argnum3, v1, v2, v3, *args, **kwargs):
+    def compute_jvp3(
+        cls, fun, argnum1, argnum2, argnum3, v1, v2, v3, *args, has_aux=False, **kwargs
+    ):
         """Compute d^3f/dx^3*v1*v2*v3.
 
         Parameters
@@ -305,17 +327,21 @@ def compute_jvp3(cls, fun, argnum1, argnum2, argnum3, v1, v2, v3, *args, **kwarg
             argnum3 = tuple([i + 2 for i in argnum3])
             v3 = tuple(v3)
 
-        dfdx = lambda dx1, *args: cls.compute_jvp(fun, argnum1, dx1, *args, **kwargs)
+        dfdx = lambda dx1, *args: cls.compute_jvp(
+            fun, argnum1, dx1, *args, has_aux=has_aux, **kwargs
+        )
         d2fdx2 = lambda dx1, dx2, *args: cls.compute_jvp(
-            dfdx, argnum2, dx2, dx1, *args, **kwargs
+            dfdx, argnum2, dx2, dx1, *args, has_aux=has_aux, **kwargs
         )
         d3fdx3 = lambda dx1, dx2, dx3: cls.compute_jvp(
-            d2fdx2, argnum3, dx3, dx2, dx1, *args, **kwargs
+            d2fdx2, argnum3, dx3, dx2, dx1, *args, has_aux=has_aux, **kwargs
         )
         return d3fdx3(v1, v2, v3)
 
     def _compute_jvp(self, v, *args, **kwargs):
-        return self.compute_jvp(self._fun, self.argnum, v, *args, **kwargs)
+        return self.compute_jvp(
+            self._fun, self.argnum, v, *args, has_aux=self._has_aux, **kwargs
+        )
 
     def _set_mode(self, mode) -> None:
         if mode not in ["fwd", "rev", "grad", "hess", "jvp"]:
@@ -324,18 +350,30 @@ def _set_mode(self, mode) -> None:
         self._mode = mode
         if self._mode == "fwd":
             self._compute = jacfwd_chunked(
-                self._fun, self._argnum, chunk_size=self._chunk_size
+                self._fun,
+                self._argnum,
+                has_aux=self._has_aux,
+                chunk_size=self._chunk_size,
             )
         elif self._mode == "rev":
             self._compute = jacrev_chunked(
-                self._fun, self._argnum, chunk_size=self._chunk_size
+                self._fun,
+                self._argnum,
+                has_aux=self._has_aux,
+                chunk_size=self._chunk_size,
             )
         elif self._mode == "grad":
-            self._compute = jax.grad(self._fun, self._argnum)
+            self._compute = jax.grad(self._fun, self._argnum, has_aux=self._has_aux)
         elif self._mode == "hess":
             self._compute = jacfwd_chunked(
-                jacrev_chunked(self._fun, self._argnum, chunk_size=self._chunk_size),
+                jacrev_chunked(
+                    self._fun,
+                    self._argnum,
+                    has_aux=self._has_aux,
+                    chunk_size=self._chunk_size,
+                ),
                 self._argnum,
+                has_aux=self._has_aux,
                 chunk_size=self._chunk_size,
             )
         elif self._mode == "jvp":
diff --git a/desc/optimize/aug_lagrangian.py b/desc/optimize/aug_lagrangian.py
index 30a2055d11..99cac57f3e 100644
--- a/desc/optimize/aug_lagrangian.py
+++ b/desc/optimize/aug_lagrangian.py
@@ -25,6 +25,7 @@
     inequality_to_bounds,
     print_header_nonlinear,
     print_iteration_nonlinear,
+    wrap_stateless_fun,
 )
 
 
@@ -45,6 +46,10 @@ def fmin_auglag(  # noqa: C901
     maxiter=None,
     callback=None,
     options={},
+    fun_has_state=False,
+    con_has_state=False,
+    fun_init_state=None,
+    con_init_state=None,
 ):
     """Minimize a function with constraints using an augmented Lagrangian method.
 
@@ -192,6 +197,19 @@ def fmin_auglag(  # noqa: C901
         - ``"scaled_termination"`` : Whether to evaluate termination criteria for
           ``xtol`` and ``gtol`` in scaled / normalized units (default) or base units.
 
+    fun_has_state : bool
+        If True, `fun`, `grad` and `hess` are assumed to have a signature of the form
+        fun(x, *args, state=fun_state) -> array, fun_state. State will be updated on
+        each accepted step.
+    con_has_state : bool
+        If True, `constraint` is assumed to have a signature of the form
+        fun(x, *args, state=con_state) -> array, con_state. State will be updated on
+        each accepted step.
+    fun_init_state : Any
+        Initial value for fun_state. Only used if `fun_has_state=True`.
+    con_init_state : Any
+        Initial value for con_state. Only used if `con_has_state=True`.
+
     Returns
     -------
     res : OptimizeResult
@@ -207,15 +225,42 @@ def fmin_auglag(  # noqa: C901
            Second Edition (2006).
 
     """
-    constraint = setdefault(
-        constraint,
-        NonlinearConstraint(  # create a dummy constraint
+    if constraint is None:
+        constraint = NonlinearConstraint(  # create a dummy constraint
             fun=lambda x, *args: jnp.array([0.0]),
             lb=0.0,
             ub=0.0,
             jac=lambda x, *args: jnp.zeros((1, x.size)),
-        ),
-    )
+        )
+        con_has_state = False
+        con_init_state = None
+
+    if not fun_has_state:
+        fun = wrap_stateless_fun(fun)
+        grad = wrap_stateless_fun(grad)
+        if callable(hess):
+            hess = wrap_stateless_fun(hess)
+    if not con_has_state:
+        constraint_fun = wrap_stateless_fun(constraint.fun)
+        constraint_jac = wrap_stateless_fun(constraint.jac)
+        constraint_hess = (
+            wrap_stateless_fun(constraint.hess)
+            if callable(constraint.hess)
+            else constraint.hess
+        )
+        constraint_vjp = (
+            wrap_stateless_fun(constraint.vjp) if hasattr(constraint, "vjp") else None
+        )
+        constraint = NonlinearConstraint(
+            fun=constraint_fun,
+            lb=constraint.lb,
+            ub=constraint.ub,
+            jac=constraint_jac,
+            hess=constraint_hess,
+            keep_feasible=constraint.keep_feasible,
+        )
+        if constraint_vjp:
+            constraint.vjp = constraint_vjp
 
     (
         z0,
@@ -233,15 +278,17 @@ def fmin_auglag(  # noqa: C901
         constraint,
         bounds,
         *args,
+        con_state=con_init_state,
     )
 
     def lagfun(f, c, y, mu, *args):
         return f - jnp.dot(y, c) + jnp.sum(mu / 2 * c * c)
 
-    def laggrad(z, y, mu, *args):
-        c = constraint_wrapped.fun(z, *args)
-        yJ = constraint_wrapped.vjp(y - mu * c, z, *args)
-        return grad_wrapped(z, *args) - yJ
+    def laggrad(z, y, mu, *args, fun_state, con_state):
+        c, con_state = constraint_wrapped.fun(z, *args, state=con_state)
+        yJ, con_state = constraint_wrapped.vjp(y - mu * c, z, *args, state=con_state)
+        g, fun_state = grad_wrapped(z, *args, state=fun_state)
+        return g - yJ, fun_state, con_state
 
     if isinstance(hess_wrapped, str) and hess_wrapped.lower() == "bfgs":
         bfgs = True
@@ -266,22 +313,22 @@ def laggrad(z, y, mu, *args):
     elif callable(constraint_wrapped.hess) and callable(hess_wrapped):
         bfgs = False
 
-        def laghess(z, y, mu, *args):
-            c = constraint_wrapped.fun(z, *args)
-            Hf = hess_wrapped(z, *args)
-            Jc = constraint_wrapped.jac(z, *args)
-            Hc1 = constraint_wrapped.hess(z, y)
-            Hc2 = constraint_wrapped.hess(z, c * mu)
-            return Hf - Hc1 + Hc2 + jnp.dot(mu * Jc.T, Jc)
+        def laghess(z, y, mu, *args, fun_state, con_state):
+            c, con_state = constraint_wrapped.fun(z, *args, state=con_state)
+            Hf, fun_state = hess_wrapped(z, *args, state=fun_state)
+            Jc, con_state = constraint_wrapped.jac(z, *args, state=con_state)
+            Hc1, con_state = constraint_wrapped.hess(z, y, state=con_state)
+            Hc2, con_state = constraint_wrapped.hess(z, c * mu, state=con_state)
+            return Hf - Hc1 + Hc2 + jnp.dot(mu * Jc.T, Jc), fun_state, con_state
 
     elif callable(hess_wrapped):
         bfgs = False
 
-        def laghess(z, y, mu, *args):
-            H = hess_wrapped(z, *args)
-            J = constraint_wrapped.jac(z, *args)
+        def laghess(z, y, mu, *args, fun_state, con_state):
+            H, fun_state = hess_wrapped(z, *args, state=fun_state)
+            J, con_state = constraint_wrapped.jac(z, *args, state=con_state)
             # ignoring higher order derivatives of constraints for now
-            return H + jnp.dot(mu * J.T, J)
+            return H + jnp.dot(mu * J.T, J), fun_state, con_state
 
     errorif(
         not (isinstance(laghess, BFGS) or callable(laghess)),
@@ -295,8 +342,8 @@ def laghess(z, y, mu, *args):
     iteration = 0
 
     z = z0.copy()
-    f = fun_wrapped(z, *args)
-    c = constraint_wrapped.fun(z, *args)
+    f, fun_state = fun_wrapped(z, *args, state=fun_init_state)
+    c, con_state = constraint_wrapped.fun(z, *args, state=con_init_state)
     constr_violation = jnp.linalg.norm(c, ord=jnp.inf)
     nfev += 1
 
@@ -308,19 +355,19 @@ def laghess(z, y, mu, *args):
     mu = options.pop("initial_penalty_parameter", 10 * jnp.ones_like(c))
     y = options.pop("initial_multipliers", jnp.zeros_like(c))
     if y == "least_squares":  # use least squares multiplier estimates
-        _J = constraint_wrapped.jac(z, *args)
-        _g = grad_wrapped(z, *args)
+        _J, _ = constraint_wrapped.jac(z, *args, state=con_state)
+        _g, _ = grad_wrapped(z, *args, state=fun_state)
         y = jnp.linalg.lstsq(_J.T, _g)[0]
         y = jnp.nan_to_num(y, nan=0.0, posinf=0.0, neginf=0.0)
     y, mu, c = jnp.broadcast_arrays(y, mu, c)
 
     L = lagfun(f, c, y, mu)
-    g = laggrad(z, y, mu, *args)
+    g, _, _ = laggrad(z, y, mu, *args, fun_state=fun_state, con_state=con_state)
     ngev += 1
     if bfgs:
         H = laghess.get_matrix()
     else:
-        H = laghess(z, y, mu, *args)
+        H, _, _ = laghess(z, y, mu, *args, fun_state=fun_state, con_state=con_state)
         nhev += 1
 
     allx = []
@@ -505,8 +552,8 @@ def laghess(z, y, mu, *args):
             step_norm = jnp.linalg.norm(step, ord=2)
 
             z_new = make_strictly_feasible(z + step, lb, ub, rstep=0)
-            f_new = fun_wrapped(z_new, *args)
-            c_new = constraint_wrapped.fun(z_new, *args)
+            f_new, fun_state_new = fun_wrapped(z_new, *args, state=fun_state)
+            c_new, con_state_new = constraint_wrapped.fun(z_new, *args, state=con_state)
             L_new = lagfun(f_new, c_new, y, mu)
             nfev += 1
 
@@ -560,16 +607,20 @@ def laghess(z, y, mu, *args):
             allx.append(z)
             f = f_new
             c = c_new
+            fun_state = fun_state_new
+            con_state = con_state_new
             constr_violation = jnp.linalg.norm(c, ord=jnp.inf)
             L = L_new
             g_old = g
-            g = laggrad(z, y, mu, *args)
+            g, _, _ = laggrad(z, y, mu, *args, fun_state=fun_state, con_state=con_state)
             ngev += 1
             if bfgs:
                 laghess.update(z - z_old, g - g_old)
                 H = laghess.get_matrix()
             else:
-                H = laghess(z, y, mu, *args)
+                H, _, _ = laghess(
+                    z, y, mu, *args, fun_state=fun_state, con_state=con_state
+                )
                 nhev += 1
 
             if hess_scale:
@@ -592,14 +643,18 @@ def laghess(z, y, mu, *args):
                     gtolk = max(omega / (jnp.mean(mu) ** alpha_omega), gtol)
                 # if we update lagrangian params, need to recompute L and J
                 L = lagfun(f, c, y, mu)
-                g = laggrad(z, y, mu, *args)
+                g, _, _ = laggrad(
+                    z, y, mu, *args, fun_state=fun_state, con_state=con_state
+                )
                 ngev += 1
 
                 if bfgs:
                     laghess.update(z - z_old, g - g_old)
                     H = laghess.get_matrix()
                 else:
-                    H = laghess(z, y, mu, *args)
+                    H, _, _ = laghess(
+                        z, y, mu, *args, fun_state=fun_state, con_state=con_state
+                    )
                     nhev += 1
 
                 if hess_scale:
@@ -674,6 +729,8 @@ def laghess(z, y, mu, *args):
         message=message,
         active_mask=active_mask,
         constr_violation=constr_violation,
+        fun_state=fun_state,
+        con_state=con_state,
     )
     if verbose > 0:
         if result["success"]:
diff --git a/desc/optimize/aug_lagrangian_ls.py b/desc/optimize/aug_lagrangian_ls.py
index ccdae0f0cb..3e02f1c53c 100644
--- a/desc/optimize/aug_lagrangian_ls.py
+++ b/desc/optimize/aug_lagrangian_ls.py
@@ -26,6 +26,7 @@
     print_header_nonlinear,
     print_iteration_nonlinear,
     solve_triangular_regularized,
+    wrap_stateless_fun,
 )
 
 
@@ -45,6 +46,10 @@ def lsq_auglag(  # noqa: C901
     maxiter=None,
     callback=None,
     options={},
+    fun_has_state=False,
+    con_has_state=False,
+    fun_init_state=None,
+    con_init_state=None,
 ):
     """Minimize a function with constraints using an augmented Lagrangian method.
 
@@ -176,6 +181,19 @@ def lsq_auglag(  # noqa: C901
         - ``"scaled_termination"`` : Whether to evaluate termination criteria for
           ``xtol`` and ``gtol`` in scaled / normalized units (default) or base units.
 
+    fun_has_state : bool
+        If True, `fun` and `jac` are assumed to have a signature of the form
+        fun(x, *args, state=fun_state) -> array, fun_state. State will be updated on
+        each accepted step.
+    con_has_state : bool
+        If True, `constraint` is assumed to have a signature of the form
+        fun(x, *args, state=con_state) -> array, con_state. State will be updated on
+        each accepted step.
+    fun_init_state : Any
+        Initial value for fun_state. Only used if `fun_has_state=True`.
+    con_init_state : Any
+        Initial value for con_state. Only used if `con_has_state=True`.
+
     Returns
     -------
     res : OptimizeResult
@@ -189,15 +207,29 @@ def lsq_auglag(  # noqa: C901
            methods" (2000).
 
     """
-    constraint = setdefault(
-        constraint,
-        NonlinearConstraint(  # create a dummy constraint
+    if constraint is None:
+        constraint = NonlinearConstraint(  # create a dummy constraint
             fun=lambda x, *args: jnp.array([0.0]),
             lb=0.0,
             ub=0.0,
             jac=lambda x, *args: jnp.zeros((1, x.size)),
-        ),
-    )
+        )
+        con_has_state = False
+        con_init_state = None
+
+    if not fun_has_state:
+        fun = wrap_stateless_fun(fun)
+        jac = wrap_stateless_fun(jac)
+    if not con_has_state:
+        constraint_fun = wrap_stateless_fun(constraint.fun)
+        constraint_jac = wrap_stateless_fun(constraint.jac)
+        constraint = NonlinearConstraint(
+            fun=constraint_fun,
+            lb=constraint.lb,
+            ub=constraint.ub,
+            jac=constraint_jac,
+            keep_feasible=constraint.keep_feasible,
+        )
 
     (
         z0,
@@ -215,6 +247,7 @@ def lsq_auglag(  # noqa: C901
         constraint,
         bounds,
         *args,
+        con_state=con_init_state,
     )
 
     # L(x,y,mu) = 1/2 f(x)^2 - y*c(x) + mu/2 c(x)^2 + y^2/(2*mu)
@@ -225,20 +258,20 @@ def lagfun(f, c, y, mu):
         c = -y / sqrt_mu + sqrt_mu * c
         return jnp.concatenate((f, c))
 
-    def lagjac(z, y, mu, *args):
-        Jf = jac_wrapped(z, *args)
-        Jc = constraint_wrapped.jac(z, *args)
+    def lagjac(z, y, mu, *args, fun_state, con_state):
+        Jf, fun_state = jac_wrapped(z, *args, state=fun_state)
+        Jc, con_state = constraint_wrapped.jac(z, *args, state=con_state)
         Jc = jnp.sqrt(mu)[:, None] * Jc
-        return jnp.vstack((Jf, Jc))
+        return jnp.vstack((Jf, Jc)), fun_state, con_state
 
     nfev = 0
     njev = 0
     iteration = 0
 
     z = z0.copy()
-    f = fun_wrapped(z, *args)
+    f, fun_state = fun_wrapped(z, *args, state=fun_init_state)
     cost = 1 / 2 * jnp.dot(f, f)
-    c = constraint_wrapped.fun(z, *args)
+    c, con_state = constraint_wrapped.fun(z, *args, state=con_init_state)
     constr_violation = jnp.linalg.norm(c, ord=jnp.inf)
     nfev += 1
 
@@ -250,14 +283,16 @@ def lagjac(z, y, mu, *args):
     mu = options.pop("initial_penalty_parameter", 10 * jnp.ones_like(c))
     y = options.pop("initial_multipliers", jnp.zeros_like(c))
     if y == "least_squares":  # use least squares multiplier estimates
-        _J = constraint_wrapped.jac(z, *args)
-        _g = f @ jac_wrapped(z, *args)
+        _J, _ = constraint_wrapped.jac(z, *args, state=con_state)
+        _f, _ = fun_wrapped(z, *args, state=fun_state)
+        _J_f, _ = jac_wrapped(z, *args, state=fun_state)
+        _g = _f @ _J_f
         y = jnp.linalg.lstsq(_J.T, _g)[0]
         y = jnp.nan_to_num(y, nan=0.0, posinf=0.0, neginf=0.0)
     y, mu, c = jnp.broadcast_arrays(y, mu, c)
 
     L = lagfun(f, c, y, mu)
-    J = lagjac(z, y, mu, *args)
+    J, _, _ = lagjac(z, y, mu, *args, fun_state=fun_state, con_state=con_state)
     Lcost = 1 / 2 * jnp.dot(L, L)
     g = L @ J
 
@@ -462,9 +497,9 @@ def lagjac(z, y, mu, *args):
             step_norm = jnp.linalg.norm(step, ord=2)
 
             z_new = make_strictly_feasible(z + step, lb, ub, rstep=0)
-            f_new = fun_wrapped(z_new, *args)
+            f_new, fun_state_new = fun_wrapped(z_new, *args, state=fun_state)
             cost_new = 0.5 * jnp.dot(f_new, f_new)
-            c_new = constraint_wrapped.fun(z_new, *args)
+            c_new, con_state_new = constraint_wrapped.fun(z_new, *args, state=con_state)
             L_new = lagfun(f_new, c_new, y, mu)
             nfev += 1
 
@@ -518,11 +553,13 @@ def lagjac(z, y, mu, *args):
             allx.append(z)
             f = f_new
             c = c_new
+            fun_state = fun_state_new
+            con_state = con_state_new
             constr_violation = jnp.linalg.norm(c, ord=jnp.inf)
             L = L_new
             cost = cost_new
             Lcost = Lcost_new
-            J = lagjac(z, y, mu, *args)
+            J, _, _ = lagjac(z, y, mu, *args, fun_state=fun_state, con_state=con_state)
             njev += 1
             g = jnp.dot(J.T, L)
 
@@ -547,7 +584,9 @@ def lagjac(z, y, mu, *args):
                 # if we update lagrangian params, need to recompute L and J
                 L = lagfun(f, c, y, mu)
                 Lcost = 0.5 * jnp.dot(L, L)
-                J = lagjac(z, y, mu, *args)
+                J, _, _ = lagjac(
+                    z, y, mu, *args, fun_state=fun_state, con_state=con_state
+                )
                 njev += 1
                 g = jnp.dot(J.T, L)
 
@@ -622,6 +661,8 @@ def lagjac(z, y, mu, *args):
         constr_violation=constr_violation,
         allx=[z2xs(x)[0] for x in allx],
         alltr=alltr,
+        fun_state=fun_state,
+        con_state=con_state,
     )
     if verbose > 0:
         if result["success"]:
diff --git a/desc/optimize/fmin_scalar.py b/desc/optimize/fmin_scalar.py
index ee40126a93..a89e5c198d 100644
--- a/desc/optimize/fmin_scalar.py
+++ b/desc/optimize/fmin_scalar.py
@@ -24,6 +24,7 @@
     compute_hess_scale,
     print_header_nonlinear,
     print_iteration_nonlinear,
+    wrap_stateless_fun,
 )
 
 
@@ -42,6 +43,8 @@ def fmintr(  # noqa: C901
     maxiter=None,
     callback=None,
     options=None,
+    has_state=False,
+    init_state=None,
 ):
     """Minimize a scalar function using a (quasi)-Newton trust region method.
 
@@ -162,6 +165,13 @@ def fmintr(  # noqa: C901
         - ``"scaled_termination"`` : Whether to evaluate termination criteria for
           ``xtol`` and ``gtol`` in scaled / normalized units (default) or base units.
 
+    has_state : bool
+        If True, `fun`, `grad` and `hess` are assumed to have a signature of the form
+        fun(x, *args, state=state) -> array, state. State will be updated on each
+        accepted step.
+    init_state : Any
+        Initial value for state. Only used if `has_state=True`.
+
     Returns
     -------
     res : OptimizeResult
@@ -178,6 +188,12 @@ def fmintr(  # noqa: C901
 
     """
     options = {} if options is None else options
+    if not has_state:
+        fun = wrap_stateless_fun(fun)
+        grad = wrap_stateless_fun(grad)
+        if callable(hess):
+            hess = wrap_stateless_fun(hess)
+
     nfev = 0
     ngev = 0
     nhev = 0
@@ -190,9 +206,9 @@ def fmintr(  # noqa: C901
     assert in_bounds(x, lb, ub), "x0 is infeasible"
     x = make_strictly_feasible(x, lb, ub)
 
-    f = fun(x, *args)
+    f, state = fun(x, *args, state=init_state)
     nfev += 1
-    g = grad(x, *args)
+    g, _ = grad(x, *args, state=state)
     ngev += 1
 
     if isinstance(hess, str) and hess.lower() == "bfgs":
@@ -203,7 +219,7 @@ def fmintr(  # noqa: C901
         hess_min_curvature = options.pop("hessian_minimum_curvature", None)
         hess = BFGS(hess_exception_strategy, hess_min_curvature, hess_init_scale)
     if callable(hess):
-        H = hess(x, *args)
+        H, _ = hess(x, *args, state=state)
         nhev += 1
         bfgs = False
     elif isinstance(hess, BFGS):
@@ -372,7 +388,7 @@ def fmintr(  # noqa: C901
             step_norm = jnp.linalg.norm(step, ord=2)
 
             x_new = make_strictly_feasible(x + step, lb, ub, rstep=0)
-            f_new = fun(x_new, *args)
+            f_new, state_new = fun(x_new, *args, state=state)
             nfev += 1
             actual_reduction = f - f_new
 
@@ -419,14 +435,15 @@ def fmintr(  # noqa: C901
             x_old = x
             x = x_new
             f = f_new
+            state = state_new
             g_old = g
-            g = grad(x, *args)
+            g, _ = grad(x, *args, state=state)
             ngev += 1
             if bfgs:
                 hess.update(x - x_old, g - g_old)
                 H = hess.get_matrix()
             else:
-                H = hess(x, *args)
+                H, _ = hess(x, *args, state=state)
                 nhev += 1
 
             if hess_scale:
@@ -492,6 +509,7 @@ def fmintr(  # noqa: C901
         active_mask=active_mask,
         allx=allx,
         alltr=alltr,
+        state=state,
     )
     if verbose > 0:
         if result["success"]:
diff --git a/desc/optimize/least_squares.py b/desc/optimize/least_squares.py
index 2d4473600e..93b49ed9d2 100644
--- a/desc/optimize/least_squares.py
+++ b/desc/optimize/least_squares.py
@@ -25,6 +25,7 @@
     print_header_nonlinear,
     print_iteration_nonlinear,
     solve_triangular_regularized,
+    wrap_stateless_fun,
 )
 
 
@@ -42,6 +43,8 @@ def lsqtr(  # noqa: C901
     maxiter=None,
     callback=None,
     options=None,
+    has_state=False,
+    init_state=None,
 ):
     """Solve a least squares problem using a (quasi)-Newton trust region method.
 
@@ -142,6 +145,13 @@ def lsqtr(  # noqa: C901
         - ``"scaled_termination"`` : Whether to evaluate termination criteria for
           ``xtol`` and ``gtol`` in scaled / normalized units (default) or base units.
 
+    has_state : bool
+        If True, `fun` and `jac` are assumed to have a signature of the form
+        fun(x, *args, state=state) -> array, state. State will be updated on each
+        accepted step.
+    init_state : Any
+        Initial value for state. Only used if `has_state=True`.
+
     Returns
     -------
     res : OptimizeResult
@@ -163,6 +173,9 @@ def lsqtr(  # noqa: C901
         ValueError,
         "x_scale should be one of 'jac', 'auto' or array-like, got {}".format(x_scale),
     )
+    if not has_state:
+        fun = wrap_stateless_fun(fun)
+        jac = wrap_stateless_fun(jac)
 
     nfev = 0
     njev = 0
@@ -175,12 +188,13 @@ def lsqtr(  # noqa: C901
     assert in_bounds(x, lb, ub), "x0 is infeasible"
     x = make_strictly_feasible(x, lb, ub)
 
-    f = fun(x, *args)
+    f, state = fun(x, *args, state=init_state)
     nfev += 1
     cost = 0.5 * jnp.dot(f, f)
     # block is needed for jaxify util which uses jax functions inside
     # jax.pure_callback and gets stuck due to async dispatch
-    J = jac(x, *args).block_until_ready()
+    J, _ = jac(x, *args, state=state)
+    J = J.block_until_ready()
     njev += 1
     g = jnp.dot(J.T, f)
 
@@ -354,7 +368,7 @@ def lsqtr(  # noqa: C901
             step_norm = jnp.linalg.norm(step, ord=2)
 
             x_new = make_strictly_feasible(x + step, lb, ub, rstep=0)
-            f_new = fun(x_new, *args)
+            f_new, state_new = fun(x_new, *args, state=state)
             nfev += 1
 
             cost_new = 0.5 * jnp.dot(f_new, f_new)
@@ -403,8 +417,9 @@ def lsqtr(  # noqa: C901
             x = x_new
             allx.append(x)
             f = f_new
+            state = state_new
             cost = cost_new
-            J = jac(x, *args)
+            J, _ = jac(x, *args, state=state)
             njev += 1
             g = jnp.dot(J.T, f)
 
@@ -466,6 +481,7 @@ def lsqtr(  # noqa: C901
         active_mask=active_mask,
         allx=allx,
         alltr=alltr,
+        state=state,
     )
     if verbose > 0:
         if result["success"]:
diff --git a/desc/optimize/stochastic.py b/desc/optimize/stochastic.py
index bbf7c7ab68..8e39d1c6b1 100644
--- a/desc/optimize/stochastic.py
+++ b/desc/optimize/stochastic.py
@@ -1,5 +1,8 @@
 """Function for minimizing a scalar function of multiple variables."""
 
+from typing import Any, Callable
+
+import equinox as eqx
 import optax
 from scipy.optimize import OptimizeResult
 
@@ -11,9 +14,19 @@
     check_termination,
     print_header_nonlinear,
     print_iteration_nonlinear,
+    wrap_stateless_fun,
 )
 
 
+class _HashableFunctionWithState(eqx.Module):
+    fun: Callable
+    state: Any
+
+    def __call__(self, x, *args):
+        # don't return new state for line search
+        return self.fun(x, *args, state=self.state)[0]
+
+
 def sgd(  # noqa: C901
     fun,
     x0,
@@ -28,6 +41,8 @@ def sgd(  # noqa: C901
     maxiter=None,
     callback=None,
     options=None,
+    has_state=False,
+    init_state=None,
 ):
     r"""Minimize a scalar function using one of stochastic gradient descent methods.
 
@@ -109,6 +124,12 @@ def sgd(  # noqa: C901
 
         For optax optimizers, hyperparameters specific to the chosen optimizer
         must be passed via the `optax-options` key of `options` dictionary.
+    has_state : bool
+        If True, `fun`, `grad` and `hess` are assumed to have a signature of the form
+        fun(x, *args, state=state) -> array, state. State will be updated on each
+        accepted step.
+    init_state : Any
+        Initial value for state. Only used if `has_state=True`.
 
     Returns
     -------
@@ -163,17 +184,22 @@ def sgd(  # noqa: C901
         x_scale = 1.0
 
     options = {} if options is None else options
+    if not has_state:
+        fun = wrap_stateless_fun(fun)
+        grad = wrap_stateless_fun(grad)
+
     nfev = 0
     ngev = 0
     iteration = 0
 
     N = x0.size
     x = x0.copy()
-    f = fun(x, *args)
+    f, state = fun(x, *args, state=init_state)
     nfev += 1
     # Scaled state xs = x / x_scale
     # Scaled gradient df/dxs = df/dx * dx/dxs = df/dx * x_scale
-    gs = grad(x, *args) * x_scale
+    gs, _ = grad(x, *args, state=state)
+    gs *= x_scale
     ngev += 1
     # scaled and unscaled norms
     g_norm = jnp.linalg.norm(gs / x_scale, ord=2)
@@ -247,8 +273,6 @@ def sgd(  # noqa: C901
     if method != "optax-custom":
         optax_method = getattr(optax, method.replace("optax-", ""))(**method_options)
     opt_state = optax_method.init(x)
-    # necessary for linesearch
-    optax_fun = lambda xs, *args: fun(xs * x_scale, *args)
 
     while True:
         success, message = check_termination(
@@ -269,15 +293,19 @@ def sgd(  # noqa: C901
         if success is not None:
             break
 
+        # need to close over current state since optax doesn't know about state.
+        # This is fine since we only update the state after successful steps.
+        optax_fun = _HashableFunctionWithState(fun, state)
         dxs, opt_state = optax_method.update(
             gs, opt_state, x / x_scale, value=f, grad=gs, value_fn=optax_fun
         )
         dx = dxs * x_scale
         x = x + dx
-        gs = grad(x, *args) * x_scale
+        fnew, state = fun(x, *args, state=state)
+        gs, _ = grad(x, *args, state=state)
+        gs *= x_scale
         g_norm = jnp.linalg.norm(gs / x_scale, ord=2)
         step_norm = jnp.linalg.norm(dx, ord=2)
-        fnew = fun(x, *args)
         df = f - fnew
         df_norm = jnp.abs(df)
         x_norm = jnp.linalg.norm(x, ord=2)
@@ -305,6 +333,7 @@ def sgd(  # noqa: C901
         nit=iteration,
         message=message,
         allx=allx,
+        state=state,
     )
     if verbose > 0:
         if result["success"]:
diff --git a/desc/optimize/utils.py b/desc/optimize/utils.py
index ad5b4a4b3c..8e15ebe4b5 100644
--- a/desc/optimize/utils.py
+++ b/desc/optimize/utils.py
@@ -9,7 +9,16 @@
 from desc.utils import Index
 
 
-def inequality_to_bounds(x0, fun, grad, hess, constraint, bounds, *args):
+def inequality_to_bounds(  # noqa: C901
+    x0,
+    fun,
+    grad,
+    hess,
+    constraint,
+    bounds,
+    *args,
+    con_state=None,
+):
     """Convert inequality constraints to bounds using slack variables.
 
     We do this by introducing slack variables s
@@ -30,6 +39,12 @@ def inequality_to_bounds(x0, fun, grad, hess, constraint, bounds, *args):
         constraint object of both equality and inequality constraints
     bounds : tuple
         lower and upper bounds for primal variables x
+    fun_has_state : bool
+        Whether the functions fun, grad, hess return a state.
+    con_has_state : bool
+        Whether the constraint return a state.
+    con_state : Any
+        Initial state to pass to constraint.
 
     Returns
     -------
@@ -45,8 +60,13 @@ def inequality_to_bounds(x0, fun, grad, hess, constraint, bounds, *args):
         function for splitting combined variable z into primal
         and slack variables x and s
 
+
+    Notes
+    -----
+    Assumes fun, grad, hess, constraint etc all take and return state.
     """
-    c0 = constraint.fun(x0, *args)
+    c0, _ = constraint.fun(x0, *args, state=con_state)
+
     ncon = c0.size
     bounds = tuple(jnp.broadcast_to(bi, x0.shape) for bi in bounds)
     cbounds = (constraint.lb, constraint.ub)
@@ -71,75 +91,75 @@ def inequality_to_bounds(x0, fun, grad, hess, constraint, bounds, *args):
     def z2xs(z):
         return z[: len(z) - nslack], z[len(z) - nslack :]
 
-    def fun_wrapped(z, *args):
+    def fun_wrapped(z, *args, state):
         x, s = z2xs(z)
-        return fun(x, *args)
+        return fun(x, *args, state=state)
 
     if hess is None:
         # assume grad is really jac of least squares
-        def grad_wrapped(z, *args):
+        def grad_wrapped(z, *args, state):
             x, s = z2xs(z)
-            g = grad(x, *args)
-            return jnp.hstack([g, jnp.zeros((g.shape[0], nslack))])
+            g, state = grad(x, *args, state=state)
+            return jnp.hstack([g, jnp.zeros((g.shape[0], nslack))]), state
 
     else:
 
-        def grad_wrapped(z, *args):
+        def grad_wrapped(z, *args, state):
             x, s = z2xs(z)
-            g = grad(x, *args)
-            return jnp.concatenate([g, jnp.zeros(nslack)])
+            g, state = grad(x, *args, state=state)
+            return jnp.concatenate([g, jnp.zeros(nslack)]), state
 
     if callable(hess):
 
-        def hess_wrapped(z, *args):
+        def hess_wrapped(z, *args, state):
             x, s = z2xs(z)
-            H = hess(x, *args)
-            return jnp.pad(H, (0, nslack))
+            H, state = hess(x, *args, state=state)
+            return jnp.pad(H, (0, nslack)), state
 
     else:  # using BFGS
         hess_wrapped = hess
 
-    def confun_wrapped(z, *args):
+    def confun_wrapped(z, *args, state):
         x, s = z2xs(z)
-        c = constraint.fun(x, *args)
+        c, state = constraint.fun(x, *args, state=state)
         sbig = jnp.zeros(ncon)
         sbig = put(sbig, ineq_mask, s)
-        return c - sbig - target
+        return c - sbig - target, state
 
-    def conjac_wrapped(z, *args):
+    def conjac_wrapped(z, *args, state):
         x, s = z2xs(z)
-        J = constraint.jac(x, *args)
+        J, state = constraint.jac(x, *args, state=state)
         I = jnp.eye(nslack)
         Js = jnp.zeros((ncon, nslack))
         Js = put(Js, Index[ineq_mask, :], -I)
-        return jnp.hstack([J, Js])
+        return jnp.hstack([J, Js]), state
 
     if callable(constraint.hess):
 
-        def conhess_wrapped(z, y, *args):
+        def conhess_wrapped(z, y, *args, state):
             x, s = z2xs(z)
-            H = constraint.hess(x, y, *args)
-            return jnp.pad(H, (0, nslack))
+            H, state = constraint.hess(x, y, *args, state=state)
+            return jnp.pad(H, (0, nslack)), state
 
     else:  # using BFGS
         conhess_wrapped = constraint.hess
 
     if hasattr(constraint, "vjp"):
 
-        def vjp_wrapped(y, z, *args):
+        def vjp_wrapped(y, z, *args, state):
             x, s = z2xs(z)
             I = jnp.eye(nslack)
             Js = jnp.zeros((ncon, nslack))
             Js = put(Js, Index[ineq_mask, :], -I)
-            vjpx = constraint.vjp(y, x, *args)
+            vjpx, state = constraint.vjp(y, x, *args, state=state)
             vjps = jnp.dot(y, Js)
-            return jnp.concatenate([vjpx, vjps])
+            return jnp.concatenate([vjpx, vjps]), state
 
     else:
 
-        def vjp_wrapped(y, z, *args):
-            J = conjac_wrapped(z, *args)
-            return jnp.dot(y, J)
+        def vjp_wrapped(y, z, *args, state):
+            J, state = conjac_wrapped(z, *args, state=state)
+            return jnp.dot(y, J), state
 
     newcon = copy.copy(constraint)
     newcon.fun = confun_wrapped
@@ -551,3 +571,8 @@ def solve_triangular_regularized(R, b, lower=False):
     Rs = R * dri[:, None]
     b = dri * b
     return solve_triangular(Rs, b, unit_diagonal=True, lower=lower)
+
+
+def wrap_stateless_fun(fun):
+    """Wrap a stateless fun to accept and return dummy state."""
+    return lambda x, *args, state=None: (fun(x, *args), None)
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index ee4275a70f..d1dcd96097 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -143,6 +143,34 @@ def test_convex_full_hess_dogleg(self):
         )
         np.testing.assert_allclose(out["x"], SCALAR_FUN_SOLN, atol=1e-8)
 
+    @pytest.mark.unit
+    def test_convex_full_hess_dogleg_state(self):
+        """Test minimizing convex stateful test function using dogleg method."""
+        # to test state we just use state to keep track of number of times
+        # the state gets updated, which should be the number of iterations + 1
+        x0 = np.ones(2)
+
+        state_fun = lambda x, *args, state: (scalar_fun(x, *args), state + 1)
+        state_grad = Derivative(state_fun, mode="grad", has_aux=True)
+        state_hess = Derivative(state_fun, mode="hess", has_aux=True)
+
+        out = fmintr(
+            state_fun,
+            x0,
+            state_grad,
+            state_hess,
+            verbose=3,
+            x_scale="hess",
+            ftol=0,
+            xtol=0,
+            gtol=1e-10,
+            options={"tr_method": "dogleg"},
+            has_state=True,
+            init_state=0,
+        )
+        np.testing.assert_allclose(out["x"], SCALAR_FUN_SOLN, atol=1e-8)
+        assert out["state"] == out["nit"] + 1
+
     @pytest.mark.unit
     def test_convex_full_hess_subspace(self):
         """Test minimizing rosenbrock function using subspace method with full hess."""
@@ -299,6 +327,41 @@ def test_optax_convex(self, method):
         )
         np.testing.assert_allclose(out["x"], SCALAR_FUN_SOLN, atol=1e-4, rtol=1e-4)
 
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "method",
+        [
+            # some of the optax optimizers
+            "optax-adam",
+            "optax-adamax",
+            "optax-lbfgs",
+            "optax-rmsprop",
+            "optax-sgd",
+        ],
+    )
+    def test_optax_convex_state(self, method):
+        """Test minimizing stateful convex test function using optax."""
+        x0 = np.ones(2)
+
+        state_fun = lambda x, *args, state: (scalar_fun(x, *args), state + 1)
+        state_grad = Derivative(state_fun, mode="grad", has_aux=True)
+
+        out = sgd(
+            state_fun,
+            x0,
+            state_grad,
+            method=method,
+            verbose=3,
+            ftol=1e-12,
+            xtol=1e-12,
+            gtol=1e-12,
+            maxiter=2000,
+            has_state=True,
+            init_state=0,
+        )
+        np.testing.assert_allclose(out["x"], SCALAR_FUN_SOLN, atol=1e-4, rtol=1e-4)
+        assert out["state"] == out["nit"] + 1  # we eval f after accepting step
+
     @pytest.mark.unit
     def test_optax_custom(self):
         """Test custom optax optimizers work."""
@@ -387,6 +450,35 @@ def res(p):
         )
         np.testing.assert_allclose(out["x"], p)
 
+    @pytest.mark.unit
+    def test_lsqtr_exact_state(self):
+        """Test minimizing stateful least squares function using exact trust region."""
+        p = np.array([1.0, 2.0, 3.0, 4.0, 1.0, 2.0])
+        x = np.linspace(-1, 1, 100)
+        y = vector_fun(x, p)
+
+        def res(p):
+            return vector_fun(x, p) - y
+
+        state_fun = lambda x, *args, state: (res(x, *args), state + 1)
+        state_jac = Derivative(state_fun, mode="fwd", has_aux=True)
+
+        rando = default_rng(seed=0)
+        p0 = p + 0.25 * (rando.random(p.size) - 0.5)
+
+        out = lsqtr(
+            state_fun,
+            p0,
+            state_jac,
+            verbose=3,
+            x_scale=1,
+            options={},
+            has_state=True,
+            init_state=0,
+        )
+        np.testing.assert_allclose(out["x"], p)
+        assert out["state"] == out["nit"] + 1
+
 
 @pytest.mark.unit
 def test_no_iterations():
@@ -956,9 +1048,17 @@ def fun(x):
         y = vecfun(x)
         return 1 / 2 * jnp.dot(y, y)
 
+    # we use different states for the objective and constraint. The objective state
+    # counts up by 1s, constraint state counts down by 2s
+
+    state_fun = lambda x, *args, state: (fun(x, *args), state + 1)
+    state_grad = jit(Derivative(state_fun, mode="grad", has_aux=True))
+    state_hess = jit(Derivative(state_fun, mode="hess", has_aux=True))
+    state_vecfun = lambda x, *args, state: (vecfun(x, *args), state + 1)
+    state_jac = jit(Derivative(state_vecfun, mode="fwd", has_aux=True))
+
     grad = jit(Derivative(fun, mode="grad"))
     hess = jit(Derivative(fun, mode="hess"))
-    jac = jit(Derivative(vecfun, mode="fwd"))
 
     @jit
     def con(x):
@@ -970,16 +1070,29 @@ def con(x):
     conjac = jit(Derivative(con, mode="fwd"))
     conhess = jit(Derivative(lambda x, v: v @ con(x), mode="hess"))
 
+    state_con = lambda x, *args, state: (con(x, *args), state - 2)
+    state_conjac = jit(Derivative(state_con, mode="fwd"))
+    state_conhess = jit(
+        Derivative(
+            lambda x, v, state: (v @ state_con(x, state=state)[0], state),
+            mode="hess",
+            has_aux=True,
+        )
+    )
+
     constraint = NonlinearConstraint(con, -np.inf, 0, conjac, conhess)
+    state_constraint = NonlinearConstraint(
+        state_con, -np.inf, 0, state_conjac, state_conhess
+    )
     x0 = rng.random(n)
 
     out1 = fmin_auglag(
-        fun,
+        state_fun,
         x0,
-        grad,
-        hess=hess,
+        state_grad,
+        hess=state_hess,
         bounds=(-jnp.inf, jnp.inf),
-        constraint=constraint,
+        constraint=state_constraint,
         args=(),
         x_scale="auto",
         ftol=0,
@@ -989,12 +1102,18 @@ def con(x):
         verbose=3,
         maxiter=None,
         options={"initial_multipliers": "least_squares"},
+        fun_has_state=True,
+        con_has_state=True,
+        fun_init_state=0,
+        con_init_state=0,
     )
     print(out1["active_mask"])
+    assert out1["fun_state"] == out1["nit"] + 1
+    assert out1["con_state"] == -2 * (out1["nit"] + 1)
     out2 = lsq_auglag(
-        vecfun,
+        state_vecfun,
         x0,
-        jac,
+        state_jac,
         bounds=(-jnp.inf, jnp.inf),
         constraint=constraint,
         args=(),
@@ -1006,7 +1125,10 @@ def con(x):
         verbose=3,
         maxiter=None,
         options={"initial_multipliers": "least_squares", "tr_method": "cho"},
+        fun_has_state=True,
+        fun_init_state=0,
     )
+    assert out2["fun_state"] == out2["nit"] + 1
 
     out3 = minimize(
         fun,