tmdsimpy.jaxΒΆ

Module of routines that exploit JAX for JIT compilation and automatic differentiation.

Many functions, methods, and classes implemented here have non-JAX versions, but not all.

Classes

NonlinearSolverOMP([config])

Parallel nonlinear solver object that contains several functions.

Modules

tmdsimpy.jax.harmonic_utils

Subset of tmdsimpy.utils.harmonic using JAX

tmdsimpy.jax.nlforces

Nonlinear forces implemented utilizing the JAX library and automatic differentiation.