Native TensorFlow trust-region minimizers (tf-trust-exact, tf-trust-ncg, tf-trust-krylov)

rabbit/fitter.py

@@ -14,18 +14,19 @@
 from rabbit import preconditioner as precond
 from rabbit import tfhelpers as tfh
 from rabbit.bbstat.bbstat import BinByBinStat
-from rabbit.callbacks import (
-    RESTART_MIN_IMPROVEMENT,
-    FitterCallback,
-    merge_callbacks,
-)
+from rabbit.callbacks import RESTART_MIN_IMPROVEMENT, FitterCallback, merge_callbacks
 from rabbit.impacts import (
     asym_impacts,
     global_asym_impacts,
     global_impacts,
     nonprofiled_impacts,
     traditional_impacts,
 )
+from rabbit.minimizer import (
+    minimize_trust_exact,
+    minimize_trust_krylov,
+    minimize_trust_ncg,
+)
 from rabbit.snapshot import Snapshotter, snapshot_on_signal
 from rabbit.tfhelpers import edmval_cov
 
@@ -2389,6 +2390,62 @@ def scipy_hess(yval):
                 logger.info(f"  - edmval: {edmval}")
             return pc.hess_to_internal(hess.__array__())
 
+        # Native (TF) minimizer counterparts of the callbacks above. Same
+        # contract and the same internal coordinates, but the gradient and
+        # Hessian stay tf tensors: with preconditioning off they never leave
+        # the device, and the subproblem factorizes there either way.
+        def native_loss(yval):
+            pc = pc_cell[0]
+            self.x.assign(pc.to_physical(yval))
+            return float(self.loss_val())
+
+        def native_closure(yval):
+            pc = pc_cell[0]
+            self.x.assign(pc.to_physical(yval))
+            val, grad, hess = self.loss_val_grad_hess()
+            if self.diagnostics:
+                cond_number = tfh.cond_number(hess)
+                logger.info(f"  - Condition number: {cond_number}")
+                edmval = tfh.edmval(grad, hess)
+                logger.info(f"  - edmval: {edmval}")
+            if pc.enabled:
+                grad = tf.constant(pc.grad_to_internal(grad.__array__()))
+                hess = tf.constant(pc.hess_to_internal(hess.__array__()))
+            return float(val), grad, hess
+
+        def native_grad_closure(yval):
+            pc = pc_cell[0]
+            self.x.assign(pc.to_physical(yval))
+            val, grad = self.loss_val_grad()
+            if pc.enabled:
+                grad = tf.constant(pc.grad_to_internal(grad.__array__()))
+            return float(val), grad
+
+        def native_set_point(yval):
+            pc = pc_cell[0]
+            self.x.assign(pc.to_physical(yval))
+
+        def native_hessp():
+            # internal-coordinate HVP, graph-compatible: the pc transform runs
+            # inside the compiled CG loop (numpy per CG iteration would defeat
+            # the on-device solve). Rebuilt per restart along with pc.
+            pc = pc_cell[0]
+            transforms = pc.tf_transforms()
+            if transforms is None:
+
+                def hessp(v):
+                    _, _, hp = self.loss_val_grad_hessp(v)
+                    return hp
+
+            else:
+                apply_T, apply_TT = transforms
+
+                def hessp(v):
+                    _, _, hp = self.loss_val_grad_hessp(apply_T(v))
+                    return apply_TT(hp)
+
+            return hessp
+
         # scipy works in internal coordinates throughout; y = 0 at the point the
         # transform was built.
         xval = pc_cell[0].from_physical(self.x.numpy())
@@ -2451,16 +2508,48 @@ def scipy_hess(yval):
             while True:
                 cb = FitterCallback(xval, self.earlyStopping, snapshotter=snapshotter)
                 try:
-                    res = scipy.optimize.minimize(
-                        scipy_loss,
-                        xval,
-                        method=self.minimizer_method,
-                        jac=True,
-                        tol=0.0,
-                        callback=cb,
-                        options=sci_opts,
-                        **info_minimize,
-                    )
+                    if self.minimizer_method == "tf-trust-exact":
+                        res = minimize_trust_exact(
+                            native_loss,
+                            native_closure,
+                            xval,
+                            gtol=sci_opts.get("gtol", 0.0),
+                            maxiter=sci_opts.get("maxiter"),
+                            callback=cb,
+                        )
+                    elif self.minimizer_method == "tf-trust-ncg":
+                        res = minimize_trust_ncg(
+                            native_loss,
+                            native_grad_closure,
+                            native_hessp(),
+                            native_set_point,
+                            xval,
+                            gtol=sci_opts.get("gtol", 0.0),
+                            maxiter=sci_opts.get("maxiter"),
+                            callback=cb,
+                        )
+                    elif self.minimizer_method == "tf-trust-krylov":
+                        res = minimize_trust_krylov(
+                            native_loss,
+                            native_grad_closure,
+                            native_hessp(),
+                            native_set_point,
+                            xval,
+                            gtol=sci_opts.get("gtol", 0.0),
+                            maxiter=sci_opts.get("maxiter"),
+                            callback=cb,
+                        )
+                    else:
+                        res = scipy.optimize.minimize(
+                            scipy_loss,
+                            xval,
+                            method=self.minimizer_method,
+                            jac=True,
+                            tol=0.0,
+                            callback=cb,
+                            options=sci_opts,
+                            **info_minimize,
+                        )
                 except Exception as ex:
                     # minimizer could have called the loss or hessp functions with "random" values, so restore the
                     # state from the end of the last iteration before the exception

rabbit/minimizer/__init__.py

@@ -0,0 +1,3 @@
+from .base import minimize_trust_exact, minimize_trust_krylov, minimize_trust_ncg
+
+__all__ = ["minimize_trust_exact", "minimize_trust_krylov", "minimize_trust_ncg"]