problems when calling diffeqsolve inside a blackjax sampler with *shard_map*.

[mamueller]

I regularly use diffrax inside mcmc samplers (blackjax) and am very impressed by the speedup and grateful.
THANKS!!!

I usually do not have problems with vmap'ing the samplers, but this is not always the smartest choice especially with samplers like nuts.
The blackjax tutorial used to advise to use pmap which worked until the recent change to shard_map as the default tool for explicit device parallelism in jax.

Using the new minimal shard_map example from the blackjax documentation as a starting point I recreated an example that involves an ode solution in the density function.
To make sure that the (device parallel) sampling procedure works I chose an ode with an explicit solution so that we can run the sampler with and without diffrax participation.

The aim of the code is to estimate the parameters A and $\omega$ of the following model:

$y=A sin( \omega t) $

from an observation of ts and ys.

I implemented the model (the sin )once directly and once as an initial value problem with diffrax.
Both implementations yield the same results if called directly with parameters.
I also call both models from a nuts sampler shard_map'ed over several chains.
In this case only the direct model survives the shard_map ed call of the sampler and produces output on several devices as expected.
The model using diffrax produces an error of this kind:

jax._src.source_info_util.JaxStackTraceBeforeTransformation: ValueError: Closure-converted function called with different dynamic arguments to the example arguments provided:
...

(complete error message below)

An interesting observation is that if the two paramters A and $\omega$ are set constanct in the ode_model (therby making it useless ... the error disappears. This might be a hint. I left this "helpful" assignment to be uncommented for testing.

I tested the code on several cpus by runnig
export JAX_PLATFORMS="cpu";export JAX_ENABLE_X64=1; export JAX_NUM_CPU_DEVICES=16 ; python minimal_error_reproduction.py
or on one cuda gpu. The error message stays the same.

The code is the following:

import jax
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
import numpy as np

import jax
import jax.numpy as jnp
import jax.scipy.stats as stats
from diffrax import diffeqsolve, Dopri5, ODETerm, SaveAt, SubSaveAt,PIDController
import blackjax
import matplotlib.pyplot as plt

@jax.jit
def model(A,omega,times):
    return A*jnp.sin(omega*times)


def vector_field(t, y, omega): 
    x, x_dot = y
    return  (
        x_dot,
        -omega**2*x,
    )
term = ODETerm(vector_field)
solver = Dopri5()

@jax.jit
def model_ode(A,omega,times):
    # iterestingly the error message in the sampler call vanishes if we set A and omega
    # (which makes the model useless but might be a hint where the problem lies)
    #(A, omega)=(5.0, 0.01)
    t0=times[0]
    t1=times[-1]
    sol = diffeqsolve(
        term, 
        solver, 
        t0=t0, 
        t1=t1, 
        dt0=0.01, 
        y0=( 
            0,
            omega*A
        ), 
        args=omega,
        saveat=SaveAt(
            subs={ "x": SubSaveAt(
                            ts=times,
                            fn=lambda t,y,args:y[0]
                        )
            }
        ),
    )
    return sol.ys["x"]
    
num_chains=jax.device_count()


measured_times=jnp.linspace(0,4*jnp.pi,10)

default={
    "A": 5.0,
    "omega": 0.01,
}



target=model(
        **default,
        times=measured_times
)
target_ode=model_ode(
        **default,
        times=measured_times
)

# make sure that both versions of the model produce the same result:
assert(jnp.allclose(
    target,
    target_ode,
    rtol=1e-5
 ))


observed=target+0.1*jax.random.normal(
    jax.random.key(43),
    shape=target.shape
)


def make_logdensity_fn(model):
    def logdensity_fn(A, omega, observed=observed,times=measured_times):
        scale=0.1
        loc=model(A,omega,times)
        logpdf = stats.norm.logpdf(observed, loc, scale)
        return jnp.sum(logpdf)
    return logdensity_fn 


def draw_samples(logdensity_fn):
    start={
        "A": 5.0,
        "omega": 0.01,
    }
    def pytree_logdensity(x):
        return logdensity_fn(**x)
    
    
    def inference_loop(rng_key, kernel, initial_state, num_samples):
    
        @jax.jit
        def one_step(state, rng_key):
            state, _ = kernel(rng_key, state)
            return state, state
    
        keys = jax.random.split(rng_key, num_samples)
        _, states = jax.lax.scan(one_step, initial_state, keys)
    
        return states
    # for a sanity check call
    # inference_loop(
    #    rng_key=rng_key,
    #    kernel=nuts.step,
    #    initial_state=nuts.init(start),
    #    num_samples=20
    #)

    
    
    inv_mass_matrix = np.array([10,1])
    step_size = 1e-3
     
    nuts = blackjax.nuts(pytree_logdensity, step_size, inv_mass_matrix)
    
    initial_positions = { 
        k: val*jnp.ones(num_chains)
        for k,val in start.items()
    }
    
    initial_states = jax.vmap(nuts.init, in_axes=(0))(initial_positions)
    rng_key=jax.random.key(42)
    mesh = jax.make_mesh((num_chains,), ('chain',))
    sharding = NamedSharding(mesh, P('chain'))
    
    rng_key, sample_key = jax.random.split(rng_key)
    sample_keys = jax.device_put(jax.random.split(sample_key, num_chains), sharding)
    initial_states_sharded = jax.device_put(initial_states, sharding)
    
    def run_one_chain(key, state):
        result = inference_loop(
            key[0], 
            nuts.step, 
            jax.tree.map(
                lambda x: x[0], 
                state
            ), 
            2_000
        )
        # the following commented code is from a blackjax example but seems to do nothing...
        #return jax.tree.map(
        #    lambda x: x[None], 
        #    result
        #)
        return result
    
    sharded_states = jax.jit(jax.shard_map(
        run_one_chain,
        mesh=mesh,
        in_specs=(P('chain'), P('chain')),
        out_specs=P('chain'),
        check_vma=False,
    ))(sample_keys, initial_states_sharded)
    return sharded_states

# the following code is successful
sharded_states=draw_samples(make_logdensity_fn(model))
jax.debug.visualize_array_sharding(sharded_states.position["A"])

# the following code fails
sharded_states_ode=draw_samples(make_logdensity_fn(model_ode))
jax.debug.visualize_array_sharding(sharded_states_ode.position["A"])

The full trace is:
Traceback (most recent call last):
File "/home/mark.mueller/pyvenv/thermic_github/lib/python3.12/site-packages/diffrax/_integrate.py", line 1456, in diffeqsolve
final_state, aux_stats = adjoint.loop(
File "/home/mark.mueller/pyvenv/thermic_github/lib/python3.12/site-packages/equinox/_module/_prebuilt.py", line 46, in call
return self.func(self.self, *args, **kwargs)
File "/home/mark.mueller/pyvenv/thermic_github/lib/python3.12/site-packages/diffrax/_adjoint.py", line 299, in loop
final_state = self._loop(
File "/home/mark.mueller/pyvenv/thermic_github/lib/python3.12/site-packages/diffrax/_integrate.py", line 685, in loop
final_state = outer_while_loop(
File "/home/mark.mueller/pyvenv/thermic_github/lib/python3.12/site-packages/equinox/internal/_loop/loop.py", line 107, in while_loop
return checkpointed_while_loop(
File "/home/mark.mueller/pyvenv/thermic_github/lib/python3.12/site-packages/equinox/internal/loop/checkpointed.py", line 249, in checkpointed_while_loop
final_val = _checkpointed_while_loop(
jax._src.source_info_util.JaxStackTraceBeforeTransformation: ValueError: Closure-converted function called with different dynamic arguments to the example arguments provided:

Called with: (
(
(
i64[],
bool[],
bool[],
State(
y=(f64[], f64[]),
tprev=f64[],
tnext=f64[],
made_jump=bool[],
solver_state=(bool[], (f64[], f64[])),
controller_state=(i32[], i32[], f64[], f64[]),
progress_meter_state=None,
result=diffrax._solution.RESULTS,
num_steps=i64[],
num_accepted_steps=i64[],
num_rejected_steps=i64[],
save_state={
'x':
SaveState(
saveat_ts_index=i64[], ts=f64[10], ys=f64[10], save_index=i64[]
)
},
dense_ts=None,
dense_infos=None,
dense_save_index=None,
event_tprev=None,
event_tnext=None,
event_dense_info=None,
event_values=None,
event_mask=None
)
),
),
{}
)
Closure-converted with: (
(
(
i64[],
bool[],
bool[],
State(
y=(f64[], f64[]),
tprev=f64[],
tnext=f64[],
made_jump=bool[],
solver_state=(bool[], (f64[], f64[])),
controller_state=(i32[], i32[], f64[], f64[]),
progress_meter_state=None,
result=diffrax._solution.RESULTS,
num_steps=i64[],
num_accepted_steps=i64[],
num_rejected_steps=i64[],
save_state={
'x':
SaveState(
saveat_ts_index=i64[], ts=f64[10], ys=f64[10], save_index=i64[]
)
},
dense_ts=None,
dense_infos=None,
dense_save_index=None,
event_tprev=None,
event_tnext=None,
event_dense_info=None,
event_values=None,
event_mask=None
)
),
),
{}
)

The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.

---

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/home/mark.mueller/pp2322/thermic/prototypes/jax/explicit_sharding/minimal_error_reproduction.py", line 177, in
sharded_states_ode=draw_samples(make_logdensity_fn(model_ode))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/mark.mueller/pp2322/thermic/prototypes/jax/explicit_sharding/minimal_error_reproduction.py", line 164, in draw_samples
sharded_states = jax.jit(jax.shard_map(
^^^^^^^^^^^^^^^^^^^^^^
File "/home/mark.mueller/pp2322/thermic/prototypes/jax/explicit_sharding/minimal_error_reproduction.py", line 148, in run_one_chain
result = inference_loop(
^^^^^^^^^^^^^^^
File "/home/mark.mueller/pp2322/thermic/prototypes/jax/explicit_sharding/minimal_error_reproduction.py", line 115, in inference_loop
_, states = jax.lax.scan(one_step, initial_state, keys)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/mark.mueller/pp2322/thermic/prototypes/jax/explicit_sharding/minimal_error_reproduction.py", line 111, in one_step
state, _ = kernel(rng_key, state)
^^^^^^^^^^^^^^^^^^^^^^
File "/home/mark.mueller/pyvenv/thermic_github/lib/python3.12/site-packages/blackjax/mcmc/nuts.py", line 220, in step_fn
return kernel(
^^^^^^^
File "/home/mark.mueller/pyvenv/thermic_github/lib/python3.12/site-packages/blackjax/mcmc/nuts.py", line 141, in kernel
proposal, info = proposal_generator(key_integrator, integrator_state, step_size)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/mark.mueller/pyvenv/thermic_github/lib/python3.12/site-packages/blackjax/mcmc/nuts.py", line 305, in propose
expansion_state, info = expand(
^^^^^^^
File "/home/mark.mueller/pyvenv/thermic_github/lib/python3.12/site-packages/blackjax/mcmc/trajectory.py", line 611, in expand
expansion_state, (is_diverging, is_turning) = jax.lax.while_loop(
^^^^^^^^^^^^^^^^^^^
File "/home/mark.mueller/pyvenv/thermic_github/lib/python3.12/site-packages/blackjax/mcmc/trajectory.py", line 554, in expand_once
) = trajectory_integrator(
^^^^^^^^^^^^^^^^^^^^^^
File "/home/mark.mueller/pyvenv/thermic_github/lib/python3.12/site-packages/blackjax/mcmc/trajectory.py", line 262, in integrate
new_integration_state, (is_diverging, has_terminated) = jax.lax.while_loop(
^^^^^^^^^^^^^^^^^^^
File "/home/mark.mueller/pyvenv/thermic_github/lib/python3.12/site-packages/blackjax/mcmc/trajectory.py", line 216, in add_one_state
new_state = integrator(trajectory.rightmost_state, direction * step_size)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/mark.mueller/pyvenv/thermic_github/lib/python3.12/site-packages/blackjax/mcmc/integrators.py", line 122, in one_step
) = operator2(
^^^^^^^^^^
File "/home/mark.mueller/pyvenv/thermic_github/lib/python3.12/site-packages/blackjax/mcmc/integrators.py", line 172, in update
logdensity, logdensity_grad = logdensity_and_grad_fn(new_position)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/mark.mueller/pyvenv/thermic_github/lib/python3.12/site-packages/equinox/internal/_loop/checkpointed.py", line 1049, in _checkpointed_while_loop_bwd
perturb_val = jax.eval_shape(_resolve_perturb_val).value
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/mark.mueller/pyvenv/thermic_github/lib/python3.12/site-packages/equinox/internal/_loop/checkpointed.py", line 1025, in _resolve_perturb_val
jax.linearize(_to_linearize, dynamic)
File "/home/mark.mueller/pyvenv/thermic_github/lib/python3.12/site-packages/equinox/internal/_loop/checkpointed.py", line 1018, in _to_linearize
_out = _body_fun(_val)
^^^^^^^^^^^^^^^
ValueError: Closure-converted function called with different dynamic arguments to the example arguments provided:

Called with: (
(
(
i64[],
bool[],
bool[],
State(
y=(f64[], f64[]),
tprev=f64[],
tnext=f64[],
made_jump=bool[],
solver_state=(bool[], (f64[], f64[])),
controller_state=(i32[], i32[], f64[], f64[]),
progress_meter_state=None,
result=diffrax._solution.RESULTS,
num_steps=i64[],
num_accepted_steps=i64[],
num_rejected_steps=i64[],
save_state={
'x':
SaveState(
saveat_ts_index=i64[], ts=f64[10], ys=f64[10], save_index=i64[]
)
},
dense_ts=None,
dense_infos=None,
dense_save_index=None,
event_tprev=None,
event_tnext=None,
event_dense_info=None,
event_values=None,
event_mask=None
)
),
),
{}
)
Closure-converted with: (
(
(
i64[],
bool[],
bool[],
State(
y=(f64[], f64[]),
tprev=f64[],
tnext=f64[],
made_jump=bool[],
solver_state=(bool[], (f64[], f64[])),
controller_state=(i32[], i32[], f64[], f64[]),
progress_meter_state=None,
result=diffrax._solution.RESULTS,
num_steps=i64[],
num_accepted_steps=i64[],
num_rejected_steps=i64[],
save_state={
'x':
SaveState(
saveat_ts_index=i64[], ts=f64[10], ys=f64[10], save_index=i64[]
)
},
dense_ts=None,
dense_infos=None,
dense_save_index=None,
event_tprev=None,
event_tnext=None,
event_dense_info=None,
event_values=None,
event_mask=None
)
),
),
{}
)

For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

and the code

[patrick-kidger]

So I rather suspect this is essentially a spurious error, with a jax.eval_shape (or something spiritually in the same space, like eqx.filter_closure_convert) being called with different input shardings.

Generally speaking I've not been a user of sharding, so these kinds of issues aren't ones that I catch routinely. If you (or your AI) are able to identify a simple fix then I'd be happy to take a pull request. That may be here or it may be in Equinox.

[jpbrodrick89]

For reference a similar problem has been reported in lineax (patrick-kidger/lineax#185) and is also one I'm keen to fix but don't have the intuition for what is the right way to approach this (or the infra to test properly).

The sticking plaster option we found helps there is throw=False or EQX_ON_ERROR=none, not sure whether that helps in this case (I've not studied your issue). I don't think this will work for diffrax (it is very strict about static arguments not changing over a vector field call which might be what's happening here).

[mamueller]

Thanks @patrick-kidger for the quick reply!
This problem is quite important for me...
My project is close to its end and relies heavily on diffrax features like SaveAt and solvers for stiff systems so that I cant just switch to jax.experimental.ode.odeint which b.t.w. survives shard_map, but is otherwise no match of course.

I have institutional access to a (very) large cluster so sharding is a (very) important part of using it effectively (especially for mcmc based parameter estimation for systems with many parameters) .
I ran into issues like the above quite frequently with diffrax but first thought that I hadn't mastered the (admittedly quite complex) semantics of shard_map (and its error messages...) or that blackjax might be less careful than it should be...
So I built a series of simple examples starting from very different locations to get an intuition of where the problems did not originate...
Only lately I started to suspect that diffrax is not "hardened" against the shard_map use cases (which seems natural if Patrick does not use it very often as he says above).
Another reason is probably that the jax policy to explicit parallelism has changed only recently with the more or less mandatory switch to shard_map making it likely that even if there had been unit tests with pmap those would have failed as soon as the internal jax implementation changed.
The mere possibility of such an error occurring has got me by surprise since due to the functional nature of jax
I would have expected a library built on top of it to be relatively immune to internal changes w.r.t. the implementation...
Seems as if I was a bit over confident.... If it is spurious as Patrick suggests it is even more terrifying...

Now I am scared to have to touch diffrax or equinox code to achieve what I deemed to be a low hanging fruit...

So after this long lead-up my question @patrick-kidger is:
How difficult do you think it would be to get diffrax immune to the new shard_map and how would you go about it?
In particular, where would you start?
(I tried simpler solvers like euler but the problem seems to lie in code common to all solvers. My example is also still very complex involving blackjax on top of diffrax)
It would be much more helpful to be able to recreate the problem with some lines using equinox and shard_map alone...
Is there an easy test that we could write and that would fail at the moment?

[mamueller]

For reference a similar problem has been reported in lineax (patrick-kidger/lineax#185) and is also one I'm keen to fix but don't have the intuition for what is the right way to approach this (or the infra to test properly).

The sticking plaster option we found helps there is throw=False or EQX_ON_ERROR=none, not sure whether that helps in this case (I've not studied your issue). I don't think this will work for diffrax (it is very strict about static arguments not changing over a vector field call which might be what's happening here).

Thank you for pointing that out. I read through the thread and get the impression that jax code at least should be transparent to sharding, but still is not in some cases...which is a bit frightening (since I actually chose jax based on this hope...)