autopdex.utility.to_jax_function

autopdex.utility.to_jax_function(subexpr, reduced_expr)

Generation of a pure JAX function based on subexpressions and reduced return-value. Can be combined with sympys common subexpression eliminator.

Return value of the function will be an jnp.ndarray Warning: Works currently only functions with basic operations. Special functions as exp() may not be jax-conforming… Warning: arguments and function name have to be set manually. The function may have to be adjusted manually.