Wrapping transformed functions in external callbacks

[alecbetancourt]

I have a question about integrating external code with JAX. I am writing a piece of physics code and am dealing with the following situation:

• I have a set of tools that perform an iterative method that are written in JAX (running on a single GPU)
• I have an energy/objective function within this iterative scheme that returns a value and an autodiff gradient
• I want to use an external library with Numpy wrappers that interfaces with a very large piece of CPU parallel Fortran code. This code performs an iterative optimization of a input, but relies on the objective function and gradient I've defined in JAX. In general, this Fortran program exhibits some non-determinism due to stochastic methods that are used during optimization, but will always produce results that are acceptable within a known margin of error. No gradient information (and therefore no custom JVP) are needed from this external callback.

A stripped down version of my current code looks something like this:

from external_wrapper import external_optimizer_function
from jax_codebase import internal_optimizer_function

def energy_fn(coord, params):
    return jnp.sum(params * coord) # dummy value

energy_and_grad = jax.jit(jax.value_and_grad(energy_fn))

def optimize_geometry(coords, params):
    optimized_coords = external_optimizer_function(coords, params, energy_and_grad, num_external_iterations=100)
    final_energy, final_grad = energy_and_grad(optimized_coords)
    return optimized_coords, final_energy, final_grad

def optimize_with_callback(coords, params):
    return jax.pure_callback(
        optimize_geometry,
        (
            jax.ShapeDtypeStruct(coords.shape, coords.dtype), # optimized coords
            jax.ShapeDtypeStruct((), coords.dtype), # final energy
            jax.ShapeDtypeStruct(coords.shape, coords.dtype), # final grad
        ),
        coords, params,
        vmap_method='sequential'
    )

coords = jnp.zeros((num_structures, num_points, num_dimensions)) 
params = jnp.zeros((num_structures, num_params))

for i in range(num_iterations):
    coords, energies_batch, gradients_batch =  jax.vmap(optimize_with_callback)(coords, params)
    params = internal_optimizer_function(coords, energies_batch, gradients_batch)
    

My main questions/concerns are as follows:

• I'm not sure if the external program being called meets the standards of functional purity needed for jax.pure_callback(). I don't particularly care about the results from the callback I/O or any non-deterministic variance from run to run of that code. In general, the parameters provided in the callback are the only ones used for the Fortran program. I haven't run into any issues so far, but I'm not sure what guarantees can be made about the correctness of this method.
• When going from JAX -> Callback -> JAX, I would not have expected the code to work. It's unclear to me what downstream effects this has on JIT caching, tracing, etc. My benchmarking results for the real code aren't particularly ideal, but are still acceptable.

[dfm]

This all looks reasonable to me! Without knowing more about the Fortran library it's hard to confidently answer your specific questions, but everything that you're doing here looks sensible, and I wouldn't be concerned about the approach.

[alecbetancourt]

Thank you for the help! The specific library is the libdlfind wrapper (https://github.com/digital-chemistry-laboratory/libdlfind) for the DL-FIND geometry optimization library.

I think my main concern was the reentry into JAX from a pure callback. That case isn't covered in the documentation as far as I can tell (https://docs.jax.dev/en/latest/external-callbacks.html). I wasn't sure if the same JAX process is used or how things would work in terms of JIT caching for the objective function that's used inside the vmapped callback when there's so much context switching between JAX and external code.

I'll keep this in mind about non-determinism as well. The callback documentation does state that functions need only be treated as pure but I wasn't sure how much this could be relaxed in general assuming you aren't doing things like exception handling that very clearly conflict with the callback interface.

[dfm]

my main concern was the reentry into JAX from a pure callback

Good question! I've done some work to ensure that re-entrant callbacks on CPU are safe, but I can imagine that there could be issues with the GPU runtime. I must admit that I'm not sure I understand the runtime well enough to give a really confident answer here! I'd probably just try it (like you are!) to see what issues I hit in practice, but it's possible that other JAX developers will have more constructive feedback.

[alecbetancourt]

This makes sense, I'll continue along this path and report back if any issues pop up. The bulk of the main JAX code is running on the GPU, but the iterative nature of the external solver and fact that it is a CPU only code means that I'll likely be restricting my limited callback into JAX to the CPU to avoid constant transfers from host to device. Thank you again for the help and clearing up some of the mechanics of this!