Jax Sharding Issue with Axis_Types

[reubenharry]

using mesh = jax.sharding.Mesh(devices=jax.devices(),axis_names= "chains"), I get an error:

ValueError: Context mesh AbstractMesh('chains': 1, axis_types={Manual: ('chains',)}) should match the mesh of sharding AbstractMesh('chains': 1, axis_types={Auto: ('chains',)}) passed to broadcast_in_dim.

I can provide more details about exactly how I'm using the sharding, but I first wanted to check if this was easily resolved with a simple change. For example, I tried manually specifying the axis_types, but I don't know where to import Manual from.

[yashk2810]

Can you paste the full stack trace and a minimal repro? From the information you have provided, the error doesn't make sense to me.

[yashk2810]

One easy way to get around this is to NOT close over values in a shard_map. Just pass them in as arguments to the shmap and this problem will go away.

[reubenharry]

Thanks so much for your help!

Can you say more about what you mean by not closing over values in this context? I'm currently doing:

p, pscalar = PartitionSpec("chains"), PartitionSpec()
parallel_execute = shard_map(
        F, mesh=mesh, in_specs=(p, p), out_specs=(p, pscalar), check_rep=False
    )

where F is some function, that I'll try to make explicit in a minimal repro shortly.

For now, here's the stack trace with `JAX_TRACEBACK_FiLTERING='off':

=========================================================== FAILURES ============================================================
________________________________________________ LinearRegressionTest.test_emaus ________________________________________________

>   samples = self.run_emaus(
        sample_init=sample_init,
        logdensity_fn=logdensity_fn,
        transform=lambda x: x,
        ndims=2,
        key=inference_key,
        diagonal_preconditioning=True,
    )

tests/mcmc/test_sampling.py:602: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
tests/mcmc/test_sampling.py:309: in run_emaus
    info, grads_per_step, _acc_prob, final_state = emaus(
blackjax/adaptation/ensemble_mclmc.py:289: in emaus
    final_state, final_adaptation_state, info2 = run_eca(
blackjax/util.py:471: in run_eca
    final_state, final_adaptation_state, info_history = parallel_execute(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

>   final_state_all, info_history = lax.scan(step, initial_state_all, xs)
E   jax._src.source_info_util.JaxStackTraceBeforeTransformation: ValueError: Context mesh AbstractMesh('chains': 1, axis_types={Manual: ('chains',)}) should match the mesh of sharding AbstractMesh('chains': 1, axis_types={Auto: ('chains',)}) passed to broadcast_in_dim. This error occurs at source:  /Users/reubencohn-gordon/blackjax/blackjax/util.py:442:44 (run_eca.<locals>.all_steps)
E   
E   The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.
E   
E   --------------------

blackjax/util.py:442: JaxStackTraceBeforeTransformation

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

self = <tests.mcmc.test_sampling.LinearRegressionTest testMethod=test_emaus>

    def test_emaus(
        self,
    ):
        """Test the MCLMC kernel."""
    
        init_key0, init_key1, inference_key = jax.random.split(self.key, 3)
    
        x_data = jax.random.normal(init_key0, shape=(1000, 1))
        y_data = 3 * x_data + jax.random.normal(init_key1, shape=x_data.shape)
    
        logposterior_fn_ = functools.partial(
            self.regression_logprob, x=x_data, preds=y_data
        )
        logdensity_fn = lambda x: logposterior_fn_(
            coefs=x["coefs"][0], log_scale=x["log_scale"][0]
        )
    
        def sample_init(key):
            key1, key2 = jax.random.split(key)
            coefs = jax.random.uniform(key1, shape=(1,), minval=1, maxval=2)
            log_scale = jax.random.uniform(key2, shape=(1,), minval=1, maxval=2)
            return {"coefs": coefs, "log_scale": log_scale}
    
>       samples = self.run_emaus(
            sample_init=sample_init,
            logdensity_fn=logdensity_fn,
            transform=lambda x: x,
            ndims=2,
            key=inference_key,
            diagonal_preconditioning=True,
        )

tests/mcmc/test_sampling.py:602: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
tests/mcmc/test_sampling.py:309: in run_emaus
    info, grads_per_step, _acc_prob, final_state = emaus(
blackjax/adaptation/ensemble_mclmc.py:289: in emaus
    final_state, final_adaptation_state, info2 = run_eca(
blackjax/util.py:471: in run_eca
    final_state, final_adaptation_state, info_history = parallel_execute(
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/experimental/shard_map.py:196: in wrapped
    out_flat = shard_map_p.bind(
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/experimental/shard_map.py:461: in bind
    return self._true_bind(*args, **params)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/core.py:520: in _true_bind
    return self.bind_with_trace(prev_trace, args, params)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/experimental/shard_map.py:466: in bind_with_trace
    return trace.process_shard_map(shard_map_p, fun, args, **params)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/experimental/shard_map.py:838: in _shard_map_impl
    outs, out_rep = _run_shmap(fun, mesh, auto, args, in_rep, check_rep,
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/experimental/shard_map.py:854: in _run_shmap
    ans = f.call_wrapped(*in_tracers)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/linear_util.py:210: in call_wrapped
    return self.f_transformed(*args, **kwargs)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/api_util.py:284: in _argnums_partial
    return _fun(*args, **kwargs)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/api_util.py:90: in flatten_fun_nokwargs
    ans = f(*py_args)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/linear_util.py:388: in _get_result_paths_thunk
    ans = _fun(*args, **kwargs)
blackjax/util.py:442: in all_steps
    final_state_all, info_history = lax.scan(step, initial_state_all, xs)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/lax/control_flow/loops.py:331: in scan
    out = scan_p.bind(*consts, *in_flat,
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/core.py:502: in bind
    return self._true_bind(*args, **params)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/core.py:520: in _true_bind
    return self.bind_with_trace(prev_trace, args, params)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/core.py:525: in bind_with_trace
    return trace.process_primitive(self, args, params)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/experimental/shard_map.py:945: in process_primitive
    out_vals = jax.jit(f)(*in_vals)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/pjit.py:341: in cache_miss
    pgle_profiler) = _python_pjit_helper(fun, jit_info, *args, **kwargs)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/pjit.py:197: in _python_pjit_helper
    out_flat = pjit_p.bind(*args_flat, **p.params)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/core.py:502: in bind
    return self._true_bind(*args, **params)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/core.py:520: in _true_bind
    return self.bind_with_trace(prev_trace, args, params)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/core.py:525: in bind_with_trace
    return trace.process_primitive(self, args, params)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/core.py:1024: in process_primitive
    return primitive.impl(*args, **params)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/pjit.py:1738: in _pjit_call_impl
    return xc._xla.pjit(
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/pjit.py:1714: in call_impl_cache_miss
    out_flat, compiled, pgle_profiler = _pjit_call_impl_python(
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/pjit.py:1662: in _pjit_call_impl_python
    compiled = _resolve_and_lower(
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/pjit.py:1629: in _resolve_and_lower
    return _pjit_lower(
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/pjit.py:1767: in _pjit_lower
    return pxla.lower_sharding_computation(
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/profiler.py:334: in wrapper
    return func(*args, **kwargs)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py:2357: in lower_sharding_computation
    nreps, tuple_args, shape_poly_state) = _cached_lowering_to_hlo(
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py:1958: in _cached_lowering_to_hlo
    lowering_result = mlir.lower_jaxpr_to_module(
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py:1215: in lower_jaxpr_to_module
    lower_jaxpr_to_fun(
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py:1705: in lower_jaxpr_to_fun
    out_vals, tokens_out = jaxpr_subcomp(
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py:1982: in jaxpr_subcomp
    ans = lower_per_platform(rule_ctx, str(eqn.primitive),
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py:2100: in lower_per_platform
    output = kept_rules[0](ctx, *rule_args, **rule_kwargs)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/experimental/shard_map.py:751: in _shard_map_lowering
    out_nodes_, tokens_out = mlir.call_lowering(
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py:2297: in call_lowering
    func_op, output_types, effects = lower_called_computation(
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py:2279: in lower_called_computation
    func_op = _lower_jaxpr_to_fun_cached(
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py:2239: in _lower_jaxpr_to_fun_cached
    func_op = lower_jaxpr_to_fun(
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py:1705: in lower_jaxpr_to_fun
    out_vals, tokens_out = jaxpr_subcomp(
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py:1982: in jaxpr_subcomp
    ans = lower_per_platform(rule_ctx, str(eqn.primitive),
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py:2100: in lower_per_platform
    output = kept_rules[0](ctx, *rule_args, **rule_kwargs)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py:2209: in f_lowered
    jaxpr, _, consts, () = pe.trace_to_jaxpr_dynamic(wrapped_fun, ctx.avals_in)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/profiler.py:334: in wrapper
    return func(*args, **kwargs)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/interpreters/partial_eval.py:2172: in trace_to_jaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/linear_util.py:210: in call_wrapped
    return self.f_transformed(*args, **kwargs)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/linear_util.py:388: in _get_result_paths_thunk
    ans = _fun(*args, **kwargs)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/lax/control_flow/loops.py:452: in _scan_impl
    yss = _map(partial(_empty_array, (length,), None), y_avals)
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/lax/control_flow/loops.py:516: in _empty_array
    return lax.broadcast(lax.empty(aval.dtype), (*prefix, *aval.shape),
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/lax/lax.py:2254: in broadcast
    return broadcast_in_dim(operand, tuple(sizes) + np.shape(operand), dims,
/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/lax/lax.py:2277: in broadcast_in_dim
    out_sharding = canonicalize_sharding(out_sharding, 'broadcast_in_dim')
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

sharding = NamedSharding(mesh=AbstractMesh('chains': 1, axis_types={Auto: ('chains',)}), spec=PartitionSpec(None,))
api_name = 'broadcast_in_dim', check_mesh_consistency = True

    def canonicalize_sharding(sharding: NamedSharding | PartitionSpec | None,
                              api_name: str, check_mesh_consistency: bool = True
                              ) -> NamedSharding | None:
      if sharding is None:
        return sharding
      if isinstance(sharding, NamedSharding) and sharding.mesh.empty:
        return None
    
      cur_mesh = mesh_lib.get_abstract_mesh()
      if isinstance(sharding, PartitionSpec):
        if cur_mesh.empty:
          raise ValueError(
              'Using PartitionSpec when you are not under a mesh context is not'
              ' allowed. Please pass a NamedSharding instance or enter into a mesh'
              f' context via `jax.sharding.use_mesh`. Got {sharding}')
        sharding = NamedSharding(cur_mesh, sharding)
      else:
        # There are cases when you have multiple meshes set. Allow that for full
        # auto mode because of existing use cases.
        # TODO(yashkatariya): Remove this once we disallow different meshes and
        # fix the existing use cases.
        if (sharding.mesh.abstract_mesh._are_all_axes_auto and
            cur_mesh._are_all_axes_auto):
          check_mesh_consistency = False
        if (check_mesh_consistency and not cur_mesh.empty and
            sharding.mesh.abstract_mesh != cur_mesh):
>         raise ValueError(
              f'Context mesh {cur_mesh} should match the mesh of sharding'
              f' {sharding.mesh.abstract_mesh} passed to {api_name}.'
              ' This error occurs at source: '
              f' {source_info_util.summarize(source_info_util.current())}')
E         ValueError: Context mesh AbstractMesh('chains': 1, axis_types={Manual: ('chains',)}) should match the mesh of sharding AbstractMesh('chains': 1, axis_types={Auto: ('chains',)}) passed to broadcast_in_dim. This error occurs at source:  /Users/reubencohn-gordon/blackjax/blackjax/util.py:442:44 (run_eca.<locals>.all_steps)

/opt/homebrew/Caskroom/miniconda/base/envs/mclmc/lib/python3.11/site-packages/jax/_src/sharding_impls.py:1286: ValueError
======================================================= 1 failed in 7.59s =======================================================

[yashk2810]

Can you say more about what you mean by not closing over values in this context?

I mean this:

y = jnp.arange(8)

def f(x):
  return jnp.sin(x) + y

shmap(f, mesh=..., ...)(jnp.arange(8))

See how y is closed over. Don't do that. Instead pass in y as an argument to f and this error will go away. It's a workaround until you give us a minimal repro but I kind of know what the problem is.

[reubenharry]

Thanks for the clarification. Since the bug appeared in somewhat fiddly code, the reproduction I have here https://github.com/reubenharry/blackjax/blob/debug_emaus/tests/mcmc/minimal_repro_3.py isn't minimal yet, but I'm still working on it.

[reubenharry]

OK, I made it a little more self-contained. The closed over value final_state (which is hard to avoid here), is the result of one shard_map call (https://github.com/reubenharry/blackjax/blob/debug_emaus/tests/mcmc/minimal_repro_3.py#L230). final_state is then used in a second shard_map call on line 319)

[reubenharry]

I was able to get the bug to go away (although I didn't really get to the route of the problem), by adding .item() on one particular line. As per the advice above, the issue did indeed seem linked to closing over a value in the function passed to the shard_map. From my perspective this is resolved, so I'll close the issue, but I should note that I think there's still a problematic behavior (which I just managed to sidestep), as shown in the error reproduction above.

Many thanks for your help!