reporting error from parallel computing using blackjax: RWState(position='ShapedArray(float32[1])', logdensity='DIFFERENT ShapedArray(float64[]) vs. ShapedArray(float32[])').

[yyang97]

Hi blackjax team,

I just came to some issues with addressing parallel computing using blackjax, about float64 and float32. I know it is not an issue from blackjax, but I am just interested in how to modify blackjax function...

import jax
import jax.numpy as jnp
import blackjax
import joblib

jax.config.update("jax_enable_x64", True)

def normal_post(theta,xobs):
    return jnp.sum(jax.scipy.stats.norm.logpdf(xobs,loc = theta))

key = jax.random.PRNGKey(1)
subkeys = jax.random.split(key, num=10000)

N = 50
theta0 =jnp.array([0.])
xobs = jax.random.normal(subkeys[0],shape = (N,))


import blackjax
rw_kernel = blackjax.mcmc.random_walk.build_additive_step()


def rw_step(key, state, step_size,xobs):
    keys = jax.random.split(key, num=2)
    return rw_kernel(
        rng_key=keys[0],
        state=state,
        logdensity_fn=lambda position: normal_post(position,xobs),
        random_step=blackjax.mcmc.random_walk.normal(sigma=step_size)
    )


# inference loop

def mcmc_run(carry, rng_key):
    """One step of sampler for given step_size."""
    state, step_size,xobs = carry
    state, info = rw_step(key=rng_key, state=state, step_size=step_size,
                        xobs = xobs)
    return (state, step_size,xobs), (state, info)

init_state = blackjax.mcmc.random_walk.RWState(
    position=theta0,
    logdensity=normal_post(theta0,xobs)

)

n_iter = 10
step_size = jnp.array([1])
logpost_key = subkeys[0]
keys = jax.random.split(key, num=n_iter)
mcmc_init = (init_state, step_size,xobs)
state_scan, out_scan = jax.lax.scan(
    f=mcmc_run,
    init=mcmc_init,
    xs=keys,
    
)

If impmented without parallel, everything works fine...

def single_core(num_key):
    jax.config.update("jax_enable_x64", True)
    xobs = jax.random.normal(subkeys[num_key],shape = (N,))
    n_iter = 10
    step_size = jnp.array([1])
    logpost_key = subkeys[0]
    keys = jax.random.split(key, num=n_iter)
    mcmc_init = (init_state, step_size,xobs)
    state_scan, out_scan = jax.lax.scan(
    f=mcmc_run,
    init=mcmc_init,
    xs=keys,
    )
    return out_scan


numbers = [0,1]
results = joblib.Parallel(n_jobs=2)(joblib.delayed(single_core)(num) for num in numbers)

However, if using joblib package, it does report the error:
TypeError: true_fun and false_fun output must have identical types, got
RWState(position='ShapedArray(float32[1])', logdensity='DIFFERENT ShapedArray(float64[]) vs. ShapedArray(float32[])').

If use jax.config.update("jax_enable_x64", False), the parallel works well.However, if using jax.config.update("jax_enable_x64", True), it just has this error.

I am not sure whether this issue is easily addressed from the balckjax functions above. For example, can we manually set RWState as float64 or anything else...

I am looking forward to valuable comments from you guys!

[junpenglao]

Hi, for parallelization you should use the native JAX functionary like vmap, pmap, xmap or other more advance Sharding features.

However, the error looks like it's just coming from the problem that float64 set up was not correct. You need to put jax.config.update("jax_enable_x64", True) in the beginning of the code, which you already did. And make sure all the inputs are in float64 (e.g., theta0 =jnp.array([0.], dtype=np.float64)). A simple fix is to change jnp to np for array initialization, because numpy default to float64.

[yyang97]

Hi @junpenglao Thank you very much for your reply! I did try your method by setting all arrays as np.float64/jnp.float64. it still reports this error.

Do you have any idea about what ShapedArray(float32[]) represents in "logdensity='DIFFERENT ShapedArray(float64[]) vs. ShapedArray(float32[])'"? Right now I guess it comes from blackjax.mcmc.random_walk.build_additive_step() and it may not be able to be fixed...

[junpenglao]

it meant the output of the logdensity function

def normal_post(theta,xobs):
    return jnp.sum(jax.scipy.stats.norm.logpdf(xobs,loc = theta))

Returns float32 initially, and after 1 step it returns float64.
Did you try also casting xobs to np.float64?

[yyang97]

Thank you so much!! your comment indicates that the issue may come from the initial state

init_state = blackjax.mcmc.random_walk.RWState(
    position=theta0,
    logdensity=normal_post(theta0,xobs)
)

instead of mcmc runs. In fact, I added theta0 =jnp.array([0.], dtype=np.float64)) in the single_core() function, everything works well now! Thank you again!