vmec_jax.free_boundary_adjoint_controllerΒΆ

JAX-visible nonlinear-controller primitives for free-boundary adjoints.

Functions

jax_visible_accepted_nonlinear_controller_directional_check_jax(...)

AD-vs-FD check for accepted/rejected JAX-visible controllers.

jax_visible_accepted_nonlinear_controller_jax(...)

Run a JAX-visible fixed-length controller with accept/reject masks.

jax_visible_masked_nonlinear_controller_directional_check_jax(...)

AD-vs-FD check for a JAX-visible masked nonlinear controller.

jax_visible_masked_nonlinear_controller_jax(...)

Run a fixed-length JAX controller with JAX-visible convergence masking.

jax_visible_nonlinear_controller_directional_check_jax(...)

AD-vs-FD check for a fully JAX-visible nonlinear controller.

jax_visible_nonlinear_controller_jax(...[, ...])

Run a nonlinear controller loop entirely through JAX lax.scan.

jax_visible_segmented_accepted_nonlinear_controller_jax(...)

Run accepted/rejected JAX-visible controllers over static segments.

pytree_directional_derivative_check_jax(...)

Compare an exact JAX directional derivative with central differences.