Optional backends and numerical solvers

External backends are optional homogeneous-state references. They are not silently substituted for native models and are not required by normal tests.

Optional backends

torch_flash.backends

Optional validation backends.

BackendCapabilities dataclass

BackendCapabilities(
    autodiff: bool,
    gpu: bool,
    fugacity_coefficients: bool,
    exact_model: str,
)

Capabilities declared by an optional validation backend.

Attributes:

Name	Type	Description
autodiff	bool	Whether backend outputs participate in PyTorch automatic differentiation.
gpu	bool	Whether the backend can execute on accelerator tensors.
fugacity_coefficients	bool	Whether component fugacity coefficients are exposed.
exact_model	str	Model identity implemented by the backend, for matched-input verification rather than name-only comparison.

CoolPropBackend

CoolPropBackend(
    names: tuple[str, ...],
    *,
    backend: Literal["HEOS", "REFPROP"] = "HEOS",
)

CoolProp AbstractState wrapper.

This backend is intended for independent validation and frozen baselines. It is CPU-only and not differentiable. HEOS uses multifluid mixture machinery with GERG-style reducing/departure functions but is not labeled as the published GERG-2008 coefficient set. Exact GERG through REFPROP requires a licensed REFPROP installation configured by the user.

Parameters:

Name	Type	Description	Default
names	tuple[str, ...]	Canonical component names in composition-axis order.	required
backend	Literal['HEOS', 'REFPROP']	CoolProp "HEOS" or user-configured "REFPROP" backend.	'HEOS'

Notes

Inputs are copied through CPU scalars/lists. Returned tensors match the caller's dtype/device but are detached from PyTorch autograd.

Source code in src/torch_flash/backends/coolprop.py

def __init__(
    self,
    names: tuple[str, ...],
    *,
    backend: Literal["HEOS", "REFPROP"] = "HEOS",
) -> None:
    try:
        from CoolProp import CoolProp
    except ImportError as exc:
        raise ImportError(
            "CoolPropBackend requires the optional 'external' dependency"
        ) from exc
    self._cp = CoolProp
    self.names = names
    fluids = "&".join(self._ALIASES.get(name, name.title()) for name in names)
    self._state = CoolProp.AbstractState(backend, fluids)
    self.capabilities = BackendCapabilities(
        autodiff=False,
        gpu=False,
        fugacity_coefficients=True,
        exact_model="REFPROP-selected model" if backend == "REFPROP" else "CoolProp HEOS",
    )

select_z

select_z(
    temperature: Tensor,
    pressure: Tensor,
    composition: Tensor,
    phase: PhaseKind = "stable",
) -> Tensor

Return CoolProp compressibility factor.

Source code in src/torch_flash/backends/coolprop.py

def select_z(
    self,
    temperature: Tensor,
    pressure: Tensor,
    composition: Tensor,
    phase: PhaseKind = "stable",
) -> Tensor:
    """Return CoolProp compressibility factor."""
    self._update(temperature, pressure, composition, phase)
    return self._result(self._state.compressibility_factor(), temperature)

molar_volume

molar_volume(
    temperature: Tensor,
    pressure: Tensor,
    composition: Tensor,
    phase: PhaseKind = "stable",
) -> Tensor

Return CoolProp molar volume.

Source code in src/torch_flash/backends/coolprop.py

def molar_volume(
    self,
    temperature: Tensor,
    pressure: Tensor,
    composition: Tensor,
    phase: PhaseKind = "stable",
) -> Tensor:
    """Return CoolProp molar volume."""
    self._update(temperature, pressure, composition, phase)
    return self._result(1.0 / self._state.rhomolar(), temperature)

log_fugacity_coefficients

log_fugacity_coefficients(
    temperature: Tensor,
    pressure: Tensor,
    composition: Tensor,
    phase: PhaseKind = "stable",
) -> Tensor

Return CoolProp log fugacity coefficients.

Source code in src/torch_flash/backends/coolprop.py

def log_fugacity_coefficients(
    self,
    temperature: Tensor,
    pressure: Tensor,
    composition: Tensor,
    phase: PhaseKind = "stable",
) -> Tensor:
    """Return CoolProp log fugacity coefficients."""
    self._update(temperature, pressure, composition, phase)
    try:
        values = [self._state.fugacity_coefficient(index) for index in range(len(self.names))]
    except (AttributeError, ValueError) as exc:
        raise ModelCapabilityError(
            "selected CoolProp backend does not expose fugacity coefficients"
        ) from exc
    return torch.log(torch.tensor(values, dtype=temperature.dtype, device=temperature.device))

TeqpBackend

TeqpBackend(
    names: tuple[str, ...], model: Any, *, exact_model: str
)

Wrap a teqp residual model as a homogeneous-state model.

teqp is deliberately optional and all values cross a CPU/NumPy boundary, so this adapter is for validation rather than differentiable production calculations. The GERG constructor uses the published hard-coded GERG-2008 residual coefficient set shipped by teqp.

Parameters:

Name	Type	Description	Default
names	tuple[str, ...]	Canonical component names in composition-axis order.	required
model	Any	Constructed teqp residual model object.	required
exact_model	str	Explicit coefficient/model identity reported through capabilities.	required

Notes

Inputs cross a CPU/NumPy boundary. Returned tensors match the caller's dtype/device but are not differentiable.

Source code in src/torch_flash/backends/teqp.py

def __init__(self, names: tuple[str, ...], model: Any, *, exact_model: str) -> None:
    self.names = names
    self._model = model
    self.capabilities = BackendCapabilities(
        autodiff=False,
        gpu=False,
        fugacity_coefficients=True,
        exact_model=exact_model,
    )

canonical_peng_robinson classmethod

canonical_peng_robinson(
    components: ComponentSet,
) -> TeqpBackend

Construct teqp's canonical Peng-Robinson model.

Source code in src/torch_flash/backends/teqp.py

@classmethod
def canonical_peng_robinson(cls, components: ComponentSet) -> TeqpBackend:
    """Construct teqp's canonical Peng-Robinson model."""
    try:
        import teqp
    except ImportError as exc:
        raise ImportError("TeqpBackend requires the optional 'external' dependency") from exc
    model = teqp.canonical_PR(
        components.critical_temperature.detach().cpu().numpy(),
        components.critical_pressure.detach().cpu().numpy(),
        components.acentric_factor.detach().cpu().numpy(),
    )
    return cls(components.names, model, exact_model="teqp canonical Peng-Robinson")

gerg2008 classmethod

gerg2008(names: tuple[str, ...]) -> TeqpBackend

Construct teqp's exact GERG-2008 residual model.

Source code in src/torch_flash/backends/teqp.py

@classmethod
def gerg2008(cls, names: tuple[str, ...]) -> TeqpBackend:
    """Construct teqp's exact GERG-2008 residual model."""
    try:
        import teqp
    except ImportError as exc:
        raise ImportError("TeqpBackend requires the optional 'external' dependency") from exc
    gerg_names = [cls._GERG_NAMES.get(name, name) for name in names]
    model = teqp.make_model({"kind": "GERG2008resid", "model": {"names": gerg_names}})
    return cls(names, model, exact_model="GERG-2008 residual (teqp)")

eoscg_2015 classmethod

eoscg_2015(names: tuple[str, ...]) -> TeqpBackend

Construct the complete 2015 EOS-CG multiparameter mixture model.

The six-component scope is CO2, water, nitrogen, oxygen, argon, and carbon monoxide. teqp loads the pure-fluid Helmholtz equations plus the Gernert--Span binary reducing and departure parameters through its multifluid model factory and versioned CoolProp-format data files.

Source code in src/torch_flash/backends/teqp.py

@classmethod
def eoscg_2015(cls, names: tuple[str, ...]) -> TeqpBackend:
    """Construct the complete 2015 EOS-CG multiparameter mixture model.

    The six-component scope is CO2, water, nitrogen, oxygen, argon, and
    carbon monoxide.  teqp loads the pure-fluid Helmholtz equations plus
    the Gernert--Span binary reducing and departure parameters through its
    multifluid model factory and versioned CoolProp-format data files.
    """
    try:
        import teqp
    except ImportError as exc:
        raise ImportError("TeqpBackend requires the optional 'external' dependency") from exc
    try:
        component_names = [cls._EOSCG_NAMES[name] for name in names]
    except KeyError as exc:
        supported = ", ".join(cls._EOSCG_NAMES)
        raise ValueError(f"EOS-CG-2015 component must be one of: {supported}") from exc
    data_path = Path(teqp.get_datapath())
    model = teqp.build_multifluid_model(
        component_names,
        str(data_path),
        str(data_path / "dev" / "mixtures" / "mixture_binary_pairs.json"),
        departurepath=str(data_path / "dev" / "mixtures" / "mixture_departure_functions.json"),
    )
    return cls(
        names,
        model,
        exact_model="EOS-CG-2015 multiparameter mixture model (teqp)",
    )

molar_volume

molar_volume(
    temperature: Tensor,
    pressure: Tensor,
    composition: Tensor,
    phase: PhaseKind = "stable",
) -> Tensor

Return a teqp molar-volume root in m3/mol.

Source code in src/torch_flash/backends/teqp.py

def molar_volume(
    self,
    temperature: Tensor,
    pressure: Tensor,
    composition: Tensor,
    phase: PhaseKind = "stable",
) -> Tensor:
    """Return a teqp molar-volume root in m3/mol."""
    density = self._select_density(temperature, pressure, composition, phase)
    return torch.tensor(
        1.0 / density,
        dtype=temperature.dtype,
        device=temperature.device,
    )

pressure

pressure(
    temperature: Tensor,
    molar_volume: Tensor,
    composition: Tensor,
) -> Tensor

Return pressure at a prescribed homogeneous density.

Source code in src/torch_flash/backends/teqp.py

def pressure(self, temperature: Tensor, molar_volume: Tensor, composition: Tensor) -> Tensor:
    """Return pressure at a prescribed homogeneous density."""
    if temperature.ndim or molar_volume.ndim or composition.ndim != 1:
        raise ValueError("teqp adapter accepts one scalar T-volume state")
    x = normalize_composition(composition)
    value = self._pressure(
        float(temperature.detach().cpu()),
        1.0 / float(molar_volume.detach().cpu()),
        np.asarray(x.detach().cpu().numpy(), dtype=np.float64),
    )
    return torch.tensor(value, dtype=temperature.dtype, device=temperature.device)

select_z

select_z(
    temperature: Tensor,
    pressure: Tensor,
    composition: Tensor,
    phase: PhaseKind = "stable",
) -> Tensor

Return compressibility factor.

Source code in src/torch_flash/backends/teqp.py

def select_z(
    self,
    temperature: Tensor,
    pressure: Tensor,
    composition: Tensor,
    phase: PhaseKind = "stable",
) -> Tensor:
    """Return compressibility factor."""
    volume = self.molar_volume(temperature, pressure, composition, phase)
    z_factor = pressure * volume / (R * temperature)
    if not bool(torch.isfinite(z_factor)):
        raise InvalidStateError("teqp produced a non-finite compressibility factor")
    return z_factor

log_fugacity_coefficients

log_fugacity_coefficients(
    temperature: Tensor,
    pressure: Tensor,
    composition: Tensor,
    phase: PhaseKind = "stable",
) -> Tensor

Return teqp residual-model fugacity coefficients.

Source code in src/torch_flash/backends/teqp.py

def log_fugacity_coefficients(
    self,
    temperature: Tensor,
    pressure: Tensor,
    composition: Tensor,
    phase: PhaseKind = "stable",
) -> Tensor:
    """Return teqp residual-model fugacity coefficients."""
    t, _, x = self._numpy_state(temperature, pressure, composition)
    density = self._select_density(temperature, pressure, composition, phase)
    values = np.asarray(self._model.get_fugacity_coefficients(t, density * x))
    if np.any(values <= 0.0) or not np.isfinite(values).all():
        raise InvalidStateError("teqp produced invalid fugacity coefficients")
    return torch.log(torch.tensor(values, dtype=temperature.dtype, device=temperature.device))

Numerical solvers

torch_flash.solvers

Reusable autodifferentiable nonlinear solvers.

BatchedNewtonResult dataclass

BatchedNewtonResult(
    solution: Tensor,
    residual: Tensor,
    residual_norm: Tensor,
    iterations: Tensor,
    converged: Tensor,
)

Solutions and diagnostics from :func:batched_damped_newton.

Attributes:

Name	Type	Description
solution	Tensor	Final variables with shape (batch, variables).
residual	Tensor	Residual vectors evaluated at solution with the same shape.
residual_norm	Tensor	Per-system infinity norms with shape (batch,).
iterations	Tensor	Per-system outer iteration counts with shape (batch,).
converged	Tensor	Per-system convergence flags with shape (batch,).

NewtonResult dataclass

NewtonResult(
    solution: Tensor,
    residual: Tensor,
    residual_norm: Tensor,
    iterations: int,
    converged: bool,
)

Solution and diagnostics from :func:damped_newton.

Attributes:

Name	Type	Description
solution	Tensor	Final iterate.
residual	Tensor	Residual vector evaluated at solution.
residual_norm	Tensor	Maximum absolute residual component.
iterations	int	Number of outer Newton iterations executed.
converged	bool	Whether residual_norm <= tolerance.

BatchedTrustRegionResult dataclass

BatchedTrustRegionResult(
    solution: Tensor,
    objective: Tensor,
    gradient_norm: Tensor,
    minimum_hessian_eigenvalue: Tensor,
    iterations: Tensor,
    converged: Tensor,
    accepted_steps: Tensor,
    rejected_steps: Tensor,
    trust_radius: Tensor,
)

Independent dense trust-region minimizations executed in one batch.

Attributes:

Name	Type	Description
solution	Tensor	Final variables with shape (batch, variables).
objective	Tensor	Final scalar objective per independent state with shape (batch,).
gradient_norm	Tensor	Final infinity-norm objective gradient per state.
minimum_hessian_eigenvalue	Tensor	Smallest eigenvalue of each final symmetrized Hessian.
iterations	Tensor	Outer trust-region iterations performed per state.
converged	Tensor	Per-state flags requiring the gradient and local-curvature gates.
accepted_steps, rejected_steps		Per-state counts based on actual versus quadratic-model reduction.
trust_radius	Tensor	Final Euclidean trust-region radius per state.

TrustRegionResult dataclass

TrustRegionResult(
    solution: Tensor,
    objective: Tensor,
    gradient_norm: Tensor,
    minimum_hessian_eigenvalue: Tensor,
    iterations: int,
    converged: bool,
    accepted_steps: int,
    rejected_steps: int,
    trust_radius: Tensor,
)

Result and convergence diagnostics from dense trust-region minimization.

Attributes:

Name	Type	Description
solution	Tensor	Final variable vector. Its dtype, device, and autodiff graph follow initial and the supplied objective.
objective	Tensor	Scalar objective evaluated at solution.
gradient_norm	Tensor	Infinity norm of the final objective gradient.
minimum_hessian_eigenvalue	Tensor	Smallest eigenvalue of the final symmetrized Hessian. A converged local minimum requires it to be nonnegative within dtype-aware roundoff.
iterations	int	Number of outer trust-region iterations.
converged	bool	Whether gradient_norm met the requested tolerance.
accepted_steps, rejected_steps		Counts based on agreement between the exact objective reduction and its local quadratic model.
trust_radius	Tensor	Final Euclidean trust-region radius.

batched_damped_newton

batched_damped_newton(
    residual_function: Callable[..., Tensor],
    initial: Tensor,
    *residual_arguments: Tensor,
    tolerance: float = 1e-08,
    max_iterations: int = 50,
    max_line_search: int = 16,
    lower_bound: Tensor | None = None,
    upper_bound: Tensor | None = None,
) -> BatchedNewtonResult

Solve independent small nonlinear systems in one PyTorch batch.

Parameters:

Name	Type	Description	Default
residual_function	Callable[..., Tensor]	Batched mapping. Its first argument is the complete variable tensor; subsequent arguments are the complete tensors in residual_arguments. It must return a tensor matching the variable shape and must not couple distinct leading-batch rows.	required
initial	Tensor	Initial variables with shape (batch, variables).	required
*residual_arguments	Tensor	Tensors whose leading dimension equals batch. They are vectorized together with the variables.	()
tolerance	float	Absolute per-system infinity-norm convergence threshold.	1e-08
max_iterations	int	Maximum damped-Newton iterations.	50
max_line_search	int	Maximum residual-decreasing backtracking trials per iteration.	16
lower_bound	Tensor | None	Optional bounds broadcastable to initial.	None
upper_bound	Tensor | None	Optional bounds broadcastable to initial.	None

Returns:

Type	Description
BatchedNewtonResult	Per-system solutions, residuals, iteration counts, and convergence.

Raises:

Type	Description
ValueError	If shapes or numerical controls are invalid.

Notes

Exact per-system Jacobian blocks come from :func:torch.func.jacrev of each residual component summed across the independent batch. This evaluates only the block diagonal implied by independent systems rather than constructing and discarding cross-batch zero blocks. Gradients through model parameters and accepted Newton iterates are preserved. A singular direct solve falls back to a batched least-squares step. Failed systems remain explicit in the returned diagnostics.

Source code in src/torch_flash/solvers/newton.py

def batched_damped_newton(
    residual_function: Callable[..., Tensor],
    initial: Tensor,
    *residual_arguments: Tensor,
    tolerance: float = 1.0e-8,
    max_iterations: int = 50,
    max_line_search: int = 16,
    lower_bound: Tensor | None = None,
    upper_bound: Tensor | None = None,
) -> BatchedNewtonResult:
    """Solve independent small nonlinear systems in one PyTorch batch.

    Parameters
    ----------
    residual_function
        Batched mapping. Its first argument is the complete variable tensor;
        subsequent arguments are the complete tensors in
        ``residual_arguments``. It must return a tensor matching the variable
        shape and must not couple distinct leading-batch rows.
    initial
        Initial variables with shape ``(batch, variables)``.
    *residual_arguments
        Tensors whose leading dimension equals ``batch``. They are vectorized
        together with the variables.
    tolerance
        Absolute per-system infinity-norm convergence threshold.
    max_iterations
        Maximum damped-Newton iterations.
    max_line_search
        Maximum residual-decreasing backtracking trials per iteration.
    lower_bound, upper_bound
        Optional bounds broadcastable to ``initial``.

    Returns
    -------
    BatchedNewtonResult
        Per-system solutions, residuals, iteration counts, and convergence.

    Raises
    ------
    ValueError
        If shapes or numerical controls are invalid.

    Notes
    -----
    Exact per-system Jacobian blocks come from
    :func:`torch.func.jacrev` of each residual component summed across the
    independent batch. This evaluates only the block diagonal implied by
    independent systems rather than constructing and discarding cross-batch
    zero blocks. Gradients through model parameters and accepted Newton
    iterates are preserved. A singular direct solve falls back to a batched
    least-squares step. Failed systems remain explicit in the returned
    diagnostics.
    """
    if initial.ndim != 2 or initial.shape[0] == 0 or initial.shape[1] == 0:
        raise ValueError("batched Newton initial values must have shape (batch, variables)")
    if tolerance <= 0.0 or max_iterations <= 0 or max_line_search <= 0:
        raise ValueError("batched Newton controls must be positive")
    if any(
        argument.ndim == 0 or argument.shape[0] != initial.shape[0]
        for argument in residual_arguments
    ):
        raise ValueError("batched Newton residual arguments must share the leading batch size")

    variables = initial.clone()

    def broadcast_bound(bound: Tensor | None, name: str) -> Tensor | None:
        if bound is None:
            return None
        try:
            return torch.broadcast_to(bound, initial.shape)
        except RuntimeError as error:
            raise ValueError(f"batched Newton {name} bound is not broadcastable") from error

    lower = broadcast_bound(lower_bound, "lower")
    upper = broadcast_bound(upper_bound, "upper")

    def project(candidate: Tensor) -> Tensor:
        if lower is not None:
            candidate = torch.maximum(candidate, lower)
        if upper is not None:
            candidate = torch.minimum(candidate, upper)
        return candidate

    variables = project(variables)
    value = residual_function(variables, *residual_arguments)
    if value.shape != variables.shape:
        raise ValueError("batched Newton residuals must match the initial variable shape")

    batch_size = initial.shape[0]
    iterations = torch.zeros(batch_size, dtype=torch.int64, device=initial.device)
    residual_norm = value.abs().amax(dim=-1)
    converged = torch.isfinite(residual_norm) & (residual_norm <= tolerance)
    for iteration in range(1, max_iterations + 1):
        active = ~converged
        if bool(torch.all(~active).detach()):
            break

        jacobian = torch.stack(
            tuple(
                torch.func.jacrev(
                    lambda current, index=residual_index: residual_function(
                        current,
                        *residual_arguments,
                    )[:, index].sum()
                )(variables)
                for residual_index in range(initial.shape[1])
            ),
            dim=-2,
        )
        expected_jacobian_shape = (*initial.shape, initial.shape[1])
        if jacobian.shape != expected_jacobian_shape:
            raise ValueError("batched Newton Jacobian shape is inconsistent")
        direct_step, information = torch.linalg.solve_ex(
            jacobian,
            -value.unsqueeze(-1),
            check_errors=False,
        )
        direct_step = direct_step.squeeze(-1)
        least_squares_step = torch.linalg.lstsq(
            jacobian,
            -value.unsqueeze(-1),
        ).solution.squeeze(-1)
        direct_is_finite = torch.isfinite(direct_step).all(dim=-1)
        step = torch.where(
            ((information == 0) & direct_is_finite).unsqueeze(-1),
            direct_step,
            least_squares_step,
        )
        step = torch.where(
            torch.isfinite(step),
            step,
            -0.1 * torch.nan_to_num(value),
        )
        step = torch.where(active.unsqueeze(-1), step, torch.zeros_like(step))

        accepted = torch.zeros(batch_size, dtype=torch.bool, device=initial.device)
        accepted_variables = variables
        factor = torch.ones(batch_size, dtype=initial.dtype, device=initial.device)
        for _ in range(max_line_search):
            candidate = project(variables + factor.unsqueeze(-1) * step)
            candidate_value = residual_function(candidate, *residual_arguments)
            candidate_norm = candidate_value.abs().amax(dim=-1)
            improving = (
                active
                & ~accepted
                & torch.isfinite(candidate_value).all(dim=-1)
                & (candidate_norm < residual_norm)
            )
            accepted_variables = torch.where(
                improving.unsqueeze(-1),
                candidate,
                accepted_variables,
            )
            accepted = accepted | improving
            factor = torch.where(accepted, factor, 0.5 * factor)
            if bool(torch.all(accepted | ~active).detach()):
                break

        fallback_variables = project(variables + factor.unsqueeze(-1) * step)
        next_variables = torch.where(
            accepted.unsqueeze(-1),
            accepted_variables,
            fallback_variables,
        )
        variables = torch.where(active.unsqueeze(-1), next_variables, variables)
        value = residual_function(variables, *residual_arguments)
        iterations = torch.where(
            active,
            iterations.new_full((), iteration),
            iterations,
        )
        residual_norm = value.abs().amax(dim=-1)
        converged = torch.isfinite(residual_norm) & (residual_norm <= tolerance)

    residual_norm = value.abs().amax(dim=-1)
    converged = torch.isfinite(residual_norm) & (residual_norm <= tolerance)
    return BatchedNewtonResult(
        variables,
        value,
        residual_norm,
        iterations,
        converged,
    )

damped_newton

damped_newton(
    residual_function: Callable[[Tensor], Tensor],
    initial: Tensor,
    *,
    tolerance: float = 1e-08,
    max_iterations: int = 50,
    max_line_search: int = 16,
    lower_bound: Tensor | None = None,
    upper_bound: Tensor | None = None,
    jacobian_refresh_interval: int = 1,
) -> NewtonResult

Solve a square nonlinear system with autodiff and backtracking.

Parameters:

Name	Type	Description	Default
residual_function	Callable[[Tensor], Tensor]	Differentiable mapping from the variable tensor to an equally sized residual tensor.	required
initial	Tensor	Initial variable tensor; dtype and device are preserved.	required
tolerance	float	Absolute infinity-norm convergence threshold.	1e-08
max_iterations	int	Maximum Newton iterations.	50
max_line_search	int	Maximum residual-decreasing backtracking trials per iteration.	16
lower_bound	Tensor | None	Optional elementwise projection bounds broadcastable to initial.	None
upper_bound	Tensor | None	Optional elementwise projection bounds broadcastable to initial.	None
jacobian_refresh_interval	int	Recompute the exact autodiff Jacobian after this many accepted iterations. Values above one use rank-one Broyden updates between refreshes; the default preserves full Newton behavior.	1

Returns:

Type	Description
NewtonResult	Final iterate and explicit convergence diagnostics.

Notes

A singular square solve falls back to a least-squares step. Failure to converge is reported in the result and is never silently converted into a successful solve.

Source code in src/torch_flash/solvers/newton.py

def damped_newton(
    residual_function: Callable[[Tensor], Tensor],
    initial: Tensor,
    *,
    tolerance: float = 1.0e-8,
    max_iterations: int = 50,
    max_line_search: int = 16,
    lower_bound: Tensor | None = None,
    upper_bound: Tensor | None = None,
    jacobian_refresh_interval: int = 1,
) -> NewtonResult:
    """Solve a square nonlinear system with autodiff and backtracking.

    Parameters
    ----------
    residual_function
        Differentiable mapping from the variable tensor to an equally sized
        residual tensor.
    initial
        Initial variable tensor; dtype and device are preserved.
    tolerance
        Absolute infinity-norm convergence threshold.
    max_iterations
        Maximum Newton iterations.
    max_line_search
        Maximum residual-decreasing backtracking trials per iteration.
    lower_bound, upper_bound
        Optional elementwise projection bounds broadcastable to ``initial``.
    jacobian_refresh_interval
        Recompute the exact autodiff Jacobian after this many accepted
        iterations. Values above one use rank-one Broyden updates between
        refreshes; the default preserves full Newton behavior.

    Returns
    -------
    NewtonResult
        Final iterate and explicit convergence diagnostics.

    Notes
    -----
    A singular square solve falls back to a least-squares step. Failure to
    converge is reported in the result and is never silently converted into a
    successful solve.
    """
    if jacobian_refresh_interval < 1:
        raise ValueError("Jacobian refresh interval must be positive")
    variables = initial.clone()

    def project(candidate: Tensor) -> Tensor:
        if lower_bound is not None:
            candidate = torch.maximum(candidate, lower_bound)
        if upper_bound is not None:
            candidate = torch.minimum(candidate, upper_bound)
        return candidate

    variables = project(variables)
    converged = False
    value = residual_function(variables)
    jacobian: Tensor | None = None
    accepted_since_refresh = jacobian_refresh_interval
    for _iteration in range(1, max_iterations + 1):
        residual_norm = value.abs().max()
        if float(residual_norm.detach()) <= tolerance:
            converged = True
            break
        if jacobian is None or accepted_since_refresh >= jacobian_refresh_interval:
            jacobian = torch.func.jacrev(residual_function)(variables)
            accepted_since_refresh = 0
        try:
            step = torch.linalg.solve(jacobian, -value)
        except torch.linalg.LinAlgError:
            step = torch.linalg.lstsq(jacobian, -value.unsqueeze(-1)).solution.squeeze(-1)
        if not bool(torch.isfinite(step).all()):
            step = -0.1 * torch.nan_to_num(value)
        accepted = False
        factor = 1.0
        for _ in range(max_line_search):
            candidate = project(variables + factor * step)
            candidate_value = residual_function(candidate)
            if bool(torch.isfinite(candidate_value.detach()).all()) and float(
                candidate_value.detach().abs().max()
            ) < float(residual_norm.detach()):
                accepted_step = candidate - variables
                residual_step = candidate_value - value
                denominator = torch.dot(accepted_step, accepted_step)
                if bool(
                    torch.isfinite(denominator)
                    & (denominator > torch.finfo(denominator.dtype).tiny)
                ):
                    jacobian = (
                        jacobian
                        + torch.outer(
                            residual_step - jacobian @ accepted_step,
                            accepted_step,
                        )
                        / denominator
                    )
                    accepted_since_refresh += 1
                else:
                    jacobian = None
                    accepted_since_refresh = jacobian_refresh_interval
                variables = candidate
                value = candidate_value
                accepted = True
                break
            factor *= 0.5
        if not accepted:
            variables = project(variables + factor * step)
            value = residual_function(variables)
            jacobian = None
            accepted_since_refresh = jacobian_refresh_interval
    residual_norm = value.abs().max()
    if float(residual_norm.detach()) <= tolerance:
        converged = True
    return NewtonResult(variables, value, residual_norm, _iteration, converged)

minimize_batched_dense_trust_region

minimize_batched_dense_trust_region(
    objective: Callable[..., Tensor],
    initial: Tensor,
    *objective_arguments: Tensor,
    is_feasible: Callable[[Tensor], Tensor] | None = None,
    initial_radius: float | Tensor = 0.5,
    maximum_radius: float | Tensor = 5.0,
    gradient_tolerance: float = 1e-09,
    max_iterations: int = 100,
    acceptance_threshold: float = 0.0,
    subproblem_tolerance: float = 1e-08,
    max_subproblem_iterations: int = 80,
) -> BatchedTrustRegionResult

Minimize independent small dense objectives in one PyTorch batch.

Parameters:

Name	Type	Description	Default
objective	Callable[..., Tensor]	Mapping from the complete variable tensor and optional complete objective_arguments to one scalar objective per batch row. Rows must be mathematically independent.	required
initial	Tensor	Finite initial variables with shape (batch, variables).	required
*objective_arguments	Tensor	Tensors sharing the leading batch size.	()
is_feasible	Callable[[Tensor], Tensor] | None	Optional predicate returning one Boolean per batch row. Infeasible trial rows are rejected without evaluating their objective. Every initial row must be feasible.	None
initial_radius	float | Tensor	Positive scalars or per-state tensors with shape (batch,).	0.5
maximum_radius	float | Tensor	Positive scalars or per-state tensors with shape (batch,).	0.5
gradient_tolerance	float		1e-09
max_iterations	float		1e-09
acceptance_threshold	float		1e-09
subproblem_tolerance	float	Controls with the same meaning as :func:minimize_dense_trust_region.	1e-08
max_subproblem_iterations	float	Controls with the same meaning as :func:minimize_dense_trust_region.	1e-08

Returns:

Type	Description
BatchedTrustRegionResult	Per-state solutions and convergence, curvature, iteration, step, and radius diagnostics.

Raises:

Type	Description
ValueError	If batch shapes, radii, controls, or objective values are invalid.

Notes

PyTorch evaluates the objective, gradients, and exact block-diagonal Hessians over the complete state batch. Only the tiny Moré-Sorensen spectral subproblem is dispatched row by row. This avoids coupling independent phase states while keeping the thermodynamic kernels vectorized on CPU or accelerator devices.

Source code in src/torch_flash/solvers/trust_region.py

def minimize_batched_dense_trust_region(
    objective: Callable[..., Tensor],
    initial: Tensor,
    *objective_arguments: Tensor,
    is_feasible: Callable[[Tensor], Tensor] | None = None,
    initial_radius: float | Tensor = 0.5,
    maximum_radius: float | Tensor = 5.0,
    gradient_tolerance: float = 1.0e-9,
    max_iterations: int = 100,
    acceptance_threshold: float = 0.0,
    subproblem_tolerance: float = 1.0e-8,
    max_subproblem_iterations: int = 80,
) -> BatchedTrustRegionResult:
    """Minimize independent small dense objectives in one PyTorch batch.

    Parameters
    ----------
    objective
        Mapping from the complete variable tensor and optional complete
        ``objective_arguments`` to one scalar objective per batch row. Rows
        must be mathematically independent.
    initial
        Finite initial variables with shape ``(batch, variables)``.
    *objective_arguments
        Tensors sharing the leading batch size.
    is_feasible
        Optional predicate returning one Boolean per batch row. Infeasible
        trial rows are rejected without evaluating their objective. Every
        initial row must be feasible.
    initial_radius, maximum_radius
        Positive scalars or per-state tensors with shape ``(batch,)``.
    gradient_tolerance, max_iterations, acceptance_threshold,
    subproblem_tolerance, max_subproblem_iterations
        Controls with the same meaning as
        :func:`minimize_dense_trust_region`.

    Returns
    -------
    BatchedTrustRegionResult
        Per-state solutions and convergence, curvature, iteration, step, and
        radius diagnostics.

    Raises
    ------
    ValueError
        If batch shapes, radii, controls, or objective values are invalid.

    Notes
    -----
    PyTorch evaluates the objective, gradients, and exact block-diagonal
    Hessians over the complete state batch. Only the tiny Moré-Sorensen
    spectral subproblem is dispatched row by row. This avoids coupling
    independent phase states while keeping the thermodynamic kernels
    vectorized on CPU or accelerator devices.
    """
    if (
        initial.ndim != 2
        or initial.shape[0] == 0
        or initial.shape[1] == 0
        or not bool(torch.isfinite(initial).all())
    ):
        raise ValueError("batched trust-region initial values must have shape (batch, variables)")
    if any(
        argument.ndim == 0 or argument.shape[0] != initial.shape[0]
        for argument in objective_arguments
    ):
        raise ValueError("batched trust-region arguments must share the leading batch size")
    if (
        gradient_tolerance <= 0.0
        or max_iterations <= 0
        or not 0.0 <= acceptance_threshold < 0.25
        or subproblem_tolerance <= 0.0
        or max_subproblem_iterations <= 0
    ):
        raise ValueError("batched trust-region controls are invalid")

    batch_size, variable_count = initial.shape
    variables = initial.clone()

    def feasible(current: Tensor) -> Tensor:
        if is_feasible is None:
            return torch.ones(
                batch_size,
                dtype=torch.bool,
                device=initial.device,
            )
        result = is_feasible(current)
        if result.shape != (batch_size,) or result.dtype != torch.bool:
            raise ValueError(
                "batched trust-region feasibility predicate must return one Boolean per batch row"
            )
        return result

    if not bool(feasible(variables).all()):
        raise ValueError("batched trust-region initial values must be feasible")

    def batch_scalar(value: float | Tensor, name: str) -> Tensor:
        tensor = torch.as_tensor(value, dtype=initial.dtype, device=initial.device)
        try:
            tensor = torch.broadcast_to(tensor, (batch_size,)).clone()
        except RuntimeError as error:
            raise ValueError(f"batched trust-region {name} is not broadcastable") from error
        return tensor

    radius = batch_scalar(initial_radius, "initial radius")
    maximum = batch_scalar(maximum_radius, "maximum radius")
    if not bool(
        torch.isfinite(radius).all()
        & torch.isfinite(maximum).all()
        & (radius > 0.0).all()
        & (maximum >= radius).all()
    ):
        raise ValueError("batched trust-region radii must be finite and positive")

    def values(current: Tensor) -> Tensor:
        result = objective(current, *objective_arguments)
        if result.shape != (batch_size,):
            raise ValueError("batched trust-region objective must return shape (batch,)")
        return result

    def derivatives(current: Tensor) -> tuple[Tensor, Tensor]:
        def total(candidate: Tensor) -> Tensor:
            return values(candidate).sum()

        gradient_function = torch.func.grad(total)
        gradient = gradient_function(current)
        hessian = torch.stack(
            tuple(
                torch.func.grad(
                    lambda candidate, index=index: gradient_function(candidate)[:, index].sum()
                )(current)
                for index in range(variable_count)
            ),
            dim=-2,
        )
        return gradient, 0.5 * (hessian + hessian.mT)

    value = values(variables)
    if not bool(torch.isfinite(value).all()):
        raise ValueError("batched trust-region objective values must be finite")
    iterations = torch.zeros(batch_size, dtype=torch.int64, device=initial.device)
    accepted_steps = torch.zeros_like(iterations)
    rejected_steps = torch.zeros_like(iterations)
    converged = torch.zeros(batch_size, dtype=torch.bool, device=initial.device)

    for iteration in range(1, max_iterations + 1):
        gradient, hessian = derivatives(variables)
        gradient_norm = gradient.abs().amax(dim=-1)
        hessian_eigenvalues = torch.linalg.eigvalsh(hessian)
        minimum_hessian_eigenvalue = hessian_eigenvalues[:, 0]
        curvature_scale = torch.maximum(
            hessian_eigenvalues.detach().abs().amax(dim=-1),
            hessian_eigenvalues.new_ones((batch_size,)),
        )
        curvature_tolerance = torch.finfo(initial.dtype).eps ** 0.5 * curvature_scale
        converged = (
            torch.isfinite(gradient_norm)
            & (gradient_norm <= gradient_tolerance)
            & (minimum_hessian_eigenvalue >= -curvature_tolerance)
        )
        active = ~converged
        if not bool(active.any().detach()):
            break

        step = torch.stack(
            tuple(
                _dense_trust_region_step(
                    gradient[index],
                    hessian[index],
                    radius[index],
                    subproblem_tolerance=subproblem_tolerance,
                    max_subproblem_iterations=max_subproblem_iterations,
                )
                if bool(active[index].detach())
                else torch.zeros_like(gradient[index])
                for index in range(batch_size)
            )
        )
        predicted_reduction = -(
            torch.einsum("bi,bi->b", gradient, step)
            + 0.5 * torch.einsum("bi,bij,bj->b", step, hessian, step)
        )
        candidate = variables + step
        candidate_is_feasible = feasible(candidate)
        safe_candidate = torch.where(
            candidate_is_feasible.unsqueeze(-1),
            candidate,
            variables,
        )
        candidate_value = torch.where(
            candidate_is_feasible,
            values(safe_candidate),
            value.new_full((), torch.inf),
        )
        actual_reduction = value - candidate_value
        valid_prediction = (
            active
            & candidate_is_feasible
            & torch.isfinite(candidate_value)
            & torch.isfinite(predicted_reduction)
            & (predicted_reduction > 0.0)
        )
        ratio = torch.where(
            valid_prediction,
            actual_reduction / predicted_reduction,
            actual_reduction.new_full((), -torch.inf),
        )
        step_norm = torch.linalg.vector_norm(step, dim=-1)
        reduction_resolution = (
            64.0
            * torch.finfo(value.dtype).eps
            * torch.maximum(value.detach().abs(), value.new_ones((batch_size,)))
        )
        roundoff_candidate = (
            valid_prediction
            & (predicted_reduction <= reduction_resolution)
            & (actual_reduction.abs() <= reduction_resolution)
        )
        if bool(roundoff_candidate.any().detach()):
            candidate_gradient, _ = derivatives(safe_candidate)
            roundoff_stationary_step = (
                roundoff_candidate
                & torch.isfinite(candidate_gradient).all(dim=-1)
                & (candidate_gradient.detach().abs().amax(dim=-1) < gradient_norm.detach())
            )
        else:
            roundoff_stationary_step = torch.zeros_like(active)

        shrink = active & ~roundoff_stationary_step & (~valid_prediction | (ratio < 0.25))
        grow = active & ~shrink & (ratio > 0.75) & (step_norm >= 0.99 * radius)
        radius = torch.where(
            shrink,
            torch.clamp_min(radius / 4.0, torch.finfo(radius.dtype).eps),
            torch.where(grow, torch.minimum(2.0 * radius, maximum), radius),
        )
        accepted = (
            active & valid_prediction & (roundoff_stationary_step | (ratio > acceptance_threshold))
        )
        variables = torch.where(accepted.unsqueeze(-1), candidate, variables)
        value = values(variables)
        iterations = torch.where(active, iterations.new_full((), iteration), iterations)
        accepted_steps = accepted_steps + accepted.to(dtype=torch.int64)
        rejected_steps = rejected_steps + (active & ~accepted).to(dtype=torch.int64)

    gradient, hessian = derivatives(variables)
    gradient_norm = gradient.abs().amax(dim=-1)
    hessian_eigenvalues = torch.linalg.eigvalsh(hessian)
    minimum_hessian_eigenvalue = hessian_eigenvalues[:, 0]
    curvature_scale = torch.maximum(
        hessian_eigenvalues.detach().abs().amax(dim=-1),
        hessian_eigenvalues.new_ones((batch_size,)),
    )
    curvature_tolerance = torch.finfo(initial.dtype).eps ** 0.5 * curvature_scale
    converged = (
        torch.isfinite(gradient_norm)
        & (gradient_norm <= gradient_tolerance)
        & (minimum_hessian_eigenvalue >= -curvature_tolerance)
    )
    return BatchedTrustRegionResult(
        variables,
        values(variables),
        gradient_norm,
        minimum_hessian_eigenvalue,
        iterations,
        converged,
        accepted_steps,
        rejected_steps,
        radius,
    )

minimize_dense_trust_region

minimize_dense_trust_region(
    objective: Callable[[Tensor], Tensor],
    initial: Tensor,
    *,
    is_feasible: Callable[[Tensor], bool] | None = None,
    initial_radius: float | Tensor | None = None,
    maximum_radius: float | Tensor | None = None,
    gradient_tolerance: float = 1e-09,
    max_iterations: int = 100,
    acceptance_threshold: float = 0.0,
    subproblem_tolerance: float = 1e-08,
    max_subproblem_iterations: int = 80,
) -> TrustRegionResult

Minimize a scalar objective with an exact small dense trust region.

This operation implements the restricted-step formulation used for phase stability and multiphase Gibbs minimization by M. Petitfrere and D. V. Nichita, "Robust and efficient Trust-Region based stability analysis and multiphase flash calculations", Fluid Phase Equilibria 362 (2014), 51-68, sections 3.1-3.4, doi:10.1016/j.fluid.2013.08.039. The dense quadratic subproblem follows the Moré-Sorensen eigenvalue-shift conditions. It is intended for the small systems typical of compositional phase calculations, not large sparse optimization.

Parameters:

Name	Type	Description	Default
objective	Callable[[Tensor], Tensor]	Differentiable mapping from one variable vector to one finite scalar. Gradients and exact Hessians are assembled with PyTorch autodiff.	required
initial	Tensor	One-dimensional starting vector.	required
is_feasible	Callable[[Tensor], bool] | None	Optional predicate for hard variable bounds. Infeasible trial steps are rejected without evaluating objective. The initial vector must be feasible.	None
initial_radius	float | Tensor | None	Positive initial Euclidean radius. The paper's flash choice ||initial|| / 10 is used when omitted, with a finite floor for a zero initial vector.	None
maximum_radius	float | Tensor | None	Positive maximum radius. Ten times the initial radius is used when omitted.	None
gradient_tolerance	float	Convergence threshold for the infinity norm of the objective gradient.	1e-09
max_iterations	int	Maximum accepted-or-rejected outer iterations.	100
acceptance_threshold	float	Minimum actual-to-predicted reduction ratio for accepting a step. Zero matches the phase-equilibrium algorithm in section 3.4.	0.0
subproblem_tolerance	float	Relative radius tolerance for the shifted-Hessian subproblem.	1e-08
max_subproblem_iterations	int	Maximum safeguarded scalar shift iterations.	80

Returns:

Type	Description
TrustRegionResult	Final variables, objective and gradient diagnostics, iteration and step counts, and the final trust radius.

Raises:

Type	Description
ValueError	If the initial tensor is not a finite vector, controls are invalid, or the objective is not a finite scalar.

Notes

The method is a local minimizer. It does not prove that a tangent-plane or Gibbs minimum is global; phase calculations still require multiple starts, material-balance checks, and stability tests. A negative or poorly predicted step is rejected rather than returned as convergence.

Source code in src/torch_flash/solvers/trust_region.py

def minimize_dense_trust_region(
    objective: Callable[[Tensor], Tensor],
    initial: Tensor,
    *,
    is_feasible: Callable[[Tensor], bool] | None = None,
    initial_radius: float | Tensor | None = None,
    maximum_radius: float | Tensor | None = None,
    gradient_tolerance: float = 1.0e-9,
    max_iterations: int = 100,
    acceptance_threshold: float = 0.0,
    subproblem_tolerance: float = 1.0e-8,
    max_subproblem_iterations: int = 80,
) -> TrustRegionResult:
    """Minimize a scalar objective with an exact small dense trust region.

    This operation implements the restricted-step formulation used for phase
    stability and multiphase Gibbs minimization by M. Petitfrere and
    D. V. Nichita, "Robust and efficient Trust-Region based stability
    analysis and multiphase flash calculations", *Fluid Phase Equilibria*
    362 (2014), 51-68, sections 3.1-3.4,
    doi:10.1016/j.fluid.2013.08.039. The dense quadratic subproblem follows
    the Moré-Sorensen eigenvalue-shift conditions. It is intended for the
    small systems typical of compositional phase calculations, not large
    sparse optimization.

    Parameters
    ----------
    objective
        Differentiable mapping from one variable vector to one finite scalar.
        Gradients and exact Hessians are assembled with PyTorch autodiff.
    initial
        One-dimensional starting vector.
    is_feasible
        Optional predicate for hard variable bounds. Infeasible trial steps
        are rejected without evaluating ``objective``. The initial vector
        must be feasible.
    initial_radius
        Positive initial Euclidean radius. The paper's flash choice
        ``||initial|| / 10`` is used when omitted, with a finite floor for a
        zero initial vector.
    maximum_radius
        Positive maximum radius. Ten times the initial radius is used when
        omitted.
    gradient_tolerance
        Convergence threshold for the infinity norm of the objective gradient.
    max_iterations
        Maximum accepted-or-rejected outer iterations.
    acceptance_threshold
        Minimum actual-to-predicted reduction ratio for accepting a step.
        Zero matches the phase-equilibrium algorithm in section 3.4.
    subproblem_tolerance
        Relative radius tolerance for the shifted-Hessian subproblem.
    max_subproblem_iterations
        Maximum safeguarded scalar shift iterations.

    Returns
    -------
    TrustRegionResult
        Final variables, objective and gradient diagnostics, iteration and
        step counts, and the final trust radius.

    Raises
    ------
    ValueError
        If the initial tensor is not a finite vector, controls are invalid, or
        the objective is not a finite scalar.

    Notes
    -----
    The method is a local minimizer. It does not prove that a tangent-plane or
    Gibbs minimum is global; phase calculations still require multiple starts,
    material-balance checks, and stability tests. A negative or poorly
    predicted step is rejected rather than returned as convergence.
    """
    if initial.ndim != 1 or initial.numel() == 0 or not bool(torch.isfinite(initial).all()):
        raise ValueError("dense trust-region initial values must be one finite vector")
    if is_feasible is not None and not is_feasible(initial):
        raise ValueError("dense trust-region initial values must be feasible")
    if (
        gradient_tolerance <= 0.0
        or max_iterations <= 0
        or not 0.0 <= acceptance_threshold < 0.25
        or subproblem_tolerance <= 0.0
        or max_subproblem_iterations <= 0
    ):
        raise ValueError("dense trust-region controls are invalid")

    variables = initial.clone()
    initial_norm = torch.linalg.vector_norm(variables.detach())
    default_radius = torch.clamp_min(initial_norm / 10.0, 0.1)
    radius = torch.as_tensor(
        default_radius if initial_radius is None else initial_radius,
        dtype=variables.dtype,
        device=variables.device,
    )
    maximum = torch.as_tensor(
        10.0 * radius if maximum_radius is None else maximum_radius,
        dtype=variables.dtype,
        device=variables.device,
    )
    if (
        radius.ndim != 0
        or maximum.ndim != 0
        or not bool(
            torch.isfinite(radius) & torch.isfinite(maximum) & (radius > 0.0) & (maximum >= radius)
        )
    ):
        raise ValueError("dense trust-region radii must be finite positive scalars")

    value = objective(variables)
    if value.ndim != 0 or not bool(torch.isfinite(value)):
        raise ValueError("dense trust-region objective must return one finite scalar")
    accepted_steps = 0
    rejected_steps = 0
    converged = False
    gradient = torch.func.grad(objective)(variables)
    gradient_norm = gradient.abs().max()
    minimum_hessian_eigenvalue = variables.new_tensor(torch.nan)

    for _iteration in range(1, max_iterations + 1):
        gradient = torch.func.grad(objective)(variables)
        gradient_norm = gradient.abs().max()
        hessian = torch.func.hessian(objective)(variables)
        symmetric_hessian = 0.5 * (hessian + hessian.mT)
        hessian_eigenvalues = torch.linalg.eigvalsh(symmetric_hessian)
        minimum_hessian_eigenvalue = hessian_eigenvalues[0]
        curvature_scale = torch.maximum(
            hessian_eigenvalues.detach().abs().amax(),
            hessian_eigenvalues.new_tensor(1.0),
        )
        curvature_tolerance = torch.finfo(hessian_eigenvalues.dtype).eps ** 0.5 * curvature_scale
        if bool(
            (gradient_norm.detach() <= gradient_tolerance)
            & (minimum_hessian_eigenvalue.detach() >= -curvature_tolerance.detach())
        ):
            converged = True
            break
        step = _dense_trust_region_step(
            gradient,
            hessian,
            radius,
            subproblem_tolerance=subproblem_tolerance,
            max_subproblem_iterations=max_subproblem_iterations,
        )
        predicted_reduction = -(torch.dot(gradient, step) + 0.5 * torch.dot(step, hessian @ step))
        candidate = variables + step
        candidate_is_feasible = is_feasible is None or is_feasible(candidate)
        candidate_value = (
            objective(candidate) if candidate_is_feasible else value.new_tensor(torch.inf)
        )
        actual_reduction = value - candidate_value
        valid_prediction = bool(
            candidate_is_feasible
            and torch.isfinite(candidate_value).detach()
            & torch.isfinite(predicted_reduction).detach()
            & (predicted_reduction.detach() > 0.0)
        )
        ratio = (
            actual_reduction / predicted_reduction
            if valid_prediction
            else value.new_tensor(-torch.inf)
        )
        step_norm = torch.linalg.vector_norm(step)
        reduction_resolution = (
            64.0
            * torch.finfo(value.dtype).eps
            * torch.maximum(value.detach().abs(), value.new_tensor(1.0))
        )
        roundoff_stationary_step = False
        if (
            valid_prediction
            and bool(predicted_reduction.detach() <= reduction_resolution)
            and bool(actual_reduction.detach().abs() <= reduction_resolution)
        ):
            candidate_gradient = torch.func.grad(objective)(candidate)
            roundoff_stationary_step = bool(
                torch.isfinite(candidate_gradient).all().detach()
                & (candidate_gradient.detach().abs().max() < gradient_norm.detach())
            )
        if roundoff_stationary_step:
            pass
        elif not valid_prediction or bool(ratio.detach() < 0.25):
            radius = torch.clamp_min(radius / 4.0, torch.finfo(radius.dtype).eps)
        elif bool((ratio.detach() > 0.75) & (step_norm.detach() >= 0.99 * radius.detach())):
            radius = torch.minimum(2.0 * radius, maximum)

        if valid_prediction and (
            roundoff_stationary_step or bool(ratio.detach() > acceptance_threshold)
        ):
            variables = candidate
            value = candidate_value
            accepted_steps += 1
        else:
            rejected_steps += 1
    gradient = torch.func.grad(objective)(variables)
    gradient_norm = gradient.abs().max()
    final_hessian = torch.func.hessian(objective)(variables)
    final_hessian_eigenvalues = torch.linalg.eigvalsh(0.5 * (final_hessian + final_hessian.mT))
    minimum_hessian_eigenvalue = final_hessian_eigenvalues[0]
    final_curvature_scale = torch.maximum(
        final_hessian_eigenvalues.detach().abs().amax(),
        final_hessian_eigenvalues.new_tensor(1.0),
    )
    final_curvature_tolerance = (
        torch.finfo(final_hessian_eigenvalues.dtype).eps ** 0.5 * final_curvature_scale
    )
    converged = converged or bool(
        (gradient_norm.detach() <= gradient_tolerance)
        & (minimum_hessian_eigenvalue.detach() >= -final_curvature_tolerance.detach())
    )
    return TrustRegionResult(
        variables,
        objective(variables),
        gradient_norm,
        minimum_hessian_eigenvalue,
        _iteration,
        converged,
        accepted_steps,
        rejected_steps,
        radius,
    )