How to make shard_map compatible with external vmap?

[Qazalbash]

Hi everyone,

I'm working on a likelihood function that uses jax.experimental.shard_map for memory-efficient computation across multiple CPU devices. This function is eventually called inside a library that applies jax.vmap over it — so I can't remove or control the vmap.

However, I'm hitting this error:

AssertionError: (1, 4)

The traceback shows it originates from inside shard_map when trying to infer the shape:

File ".../sharding.py", line 61, in _common_shard_shape
  assert len(partitions) == len(global_shape), (len(partitions), len(global_shape))

I'm guessing this means that under vmap, the global shape has an extra leading batch dimension that conflicts with the defined PartitionSpec.

Here's the minimal version of what I’m doing:

import os


os.environ["JAX_PLATFORMS"] = "cpu"
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"


from typing import Optional, Tuple

import jax
from jax import Array, numpy as jnp, random as jrd
from jax.experimental.shard_map import shard_map


key1, key2, key3 = jrd.split(jrd.PRNGKey(0), num=3)

N_devices = jax.local_device_count()


mesh = jax.make_mesh(
    (N_devices,),
    axis_names=("devices",),
    axis_types=(jax.sharding.AxisType.Manual,),
)

device_partitioning = jax.sharding.PartitionSpec(*mesh.axis_names)
sharding = jax.sharding.NamedSharding(mesh, device_partitioning)


batched_data: Array = jrd.uniform(key1, (N_devices * 8, 1000, 4))
batched_log_ref_priors: Array = jrd.uniform(key1, (N_devices * 8, 1000))


sharded_data: Array = jax.block_until_ready(
    jax.device_put(batched_data, sharding, may_alias=True)
)
sharded_log_ref_priors: Array = jax.block_until_ready(
    jax.device_put(batched_log_ref_priors, sharding, may_alias=True)
)


def likelihood_fn(x: Array, _: Optional[Array]) -> Array:
    mean = jax.lax.dynamic_index_in_dim(x, 0, keepdims=False)
    sigma = jax.lax.dynamic_index_in_dim(x, 1, keepdims=False)

    def log_prob_fn(data: Array, log_ref_prior: Array) -> Array:
        def single_event_fn(_: Array, input: Tuple[Array, Array]) -> Tuple[None, Array]:
            data, log_ref_prior = input

            log_prob = jnp.sum(-(((data - mean) / sigma) ** 2), axis=-1) - log_ref_prior

            log_prob_sum = jax.nn.logsumexp(
                log_prob,
                axis=-1,
                where=(~jnp.isneginf(log_prob)),
            )
            return None, log_prob_sum

        _, total_log_likelihood_collection = jax.lax.scan(
            single_event_fn,  # type: ignore
            None,
            (data, log_ref_prior),
        )

        return total_log_likelihood_collection

    partial_log_likelihood: jax.Array = jax.block_until_ready(
        shard_map(
            log_prob_fn,
            mesh=mesh,
            in_specs=(device_partitioning, device_partitioning),
            out_specs=device_partitioning,
        )(sharded_data, sharded_log_ref_priors)
    )

    total_log_likelihood = jnp.sum(partial_log_likelihood)

    log_likelihood = total_log_likelihood
    log_posterior = log_likelihood

    return log_posterior


x = jrd.uniform(key3, (10, 2), minval=1.0, maxval=2.0)


print(jax.vmap(likelihood_fn, in_axes=(0, None))(x, None))  # <- AssertionError here

My question is:
How can I make likelihood_fn compatible with an external vmap, while keeping shard_map working on the already sharded data?

Key constraints:

• The vmap is imposed by a library I’m using — I can’t remove or rewrite it.
• sharded_data and sharded_log_ref_priors are constant inputs.
• I want to use shard_map to parallelize across devices, and vmap to evaluate the function across parameter samples.

Would wrapping the likelihood_fn body with jax.named_call, or using with_sharding_constraint, help here? Or is there a better way to "freeze" the sharded arrays so the outer vmap doesn’t try to batch them?

Thanks in advance for any insights — really stuck here.

[yashk2810]

One problem is that you are creating the mesh with Manual mesh axis. If you don't do that and just pass the mesh to shard_map, it should work or give you a different error.

shard_map will switch the mesh axes for you.

[Qazalbash]

Hi thanks for the reply!

Although it worked for CPUs, but not for GPU.

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

Traceback (most recent call last):
  File "/home/muhammad.zeeshan/.conda/envs/gwkenv/bin/sage_n_pls_m_gs", line 10, in <module>
    sys.exit(main())
             ^^^^^^
  File "/home/muhammad.zeeshan/.conda/envs/gwkenv/lib/python3.12/site-packages/kokab/n_pls_m_gs/sage.py", line 522, in main
    handler.run(
  File "/home/muhammad.zeeshan/.conda/envs/gwkenv/lib/python3.12/site-packages/kokab/utils/flowMC_helper.py", line 305, in run
    sampler.sample(self.initial_position, self.data)
  File "/home/muhammad.zeeshan/.local/lib/python3.12/site-packages/flowMC/Sampler.py", line 204, in sample
    ) = strategy(
        ^^^^^^^^^
  File "/home/muhammad.zeeshan/.local/lib/python3.12/site-packages/flowMC/strategy/global_tuning.py", line 92, in __call__
    ) = local_sampler.sample(
        ^^^^^^^^^^^^^^^^^^^^^
  File "/home/muhammad.zeeshan/.local/lib/python3.12/site-packages/flowMC/proposal/MALA.py", line 161, in sample
    state = self.update_vmap(i, state)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/muhammad.zeeshan/.local/lib/python3.12/site-packages/flowMC/proposal/MALA.py", line 123, in update
    new_position, new_log_p, do_accept = self.kernel(
                                         ^^^^^^^^^^^^
  File "/home/muhammad.zeeshan/.local/lib/python3.12/site-packages/flowMC/proposal/MALA.py", line 82, in kernel
    _, (proposal, logprob, d_logprob) = jax.lax.scan(
                                        ^^^^^^^^^^^^^
  File "/home/muhammad.zeeshan/.local/lib/python3.12/site-packages/flowMC/proposal/MALA.py", line 49, in body
    this_log_prob, this_d_log = jax.value_and_grad(self.logpdf)(this_position, data)
                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/muhammad.zeeshan/.conda/envs/gwkenv/lib/python3.12/site-packages/jax/_src/sharding_impls.py", line 1290, in canonicalize_sharding
    raise ValueError(
ValueError: Context mesh AbstractMesh('devices': 4, axis_types=(Manual,)) should match the mesh of sharding AbstractMesh('devices': 4, axis_types=(Auto,)) passed to broadcast_in_dim. This error occurs at source:  /home/muhammad.zeeshan/.local/lib/python3.12/site-packages/flowMC/proposal/MALA.py:49:36 (MALA.body)

[Qazalbash]

The code shared is an MRE. The actual implementation can be seen here.

[yashk2810]

What happens if you pass mean and sigma as args to the shard_mapped function?

[Qazalbash]

This MRE is working fine as per your suggestion, both on CPU and GPU. Not the actual code that I meant to run.

[yashk2810]

The above MRE works for me without passing mean and sigma as args to shard_mapped function at jax and jaxlib nightly.

Can you please try with that?