jax-ml/jax · error · TypeError

{name} requires all arguments to have matching type. Got key

Error message

{name} requires all arguments to have matching type. Got key type: {core.typeof(key)} vs arg type: {core.typeof(a)}. Use jax.lax.pcast(..., to='varying') to make them match. If your key is less varying than arg, watch out for key-reuse problems.

What it means

Inside JAX's random kernels (shard_map/pvary handling), each non-key argument is pvary-cast up to the key's variance level, and then the key's type (its `mat`, the sparsity/variance matrix of a sharded type) must equal the argument's. If they still differ, JAX raises this TypeError telling you to align them manually with jax.lax.pcast(..., to='varying') and warning that a less-varying key implies key reuse across devices.

Source

Thrown at jax/_src/random/core.py:3783

def random_insert_pvary(name, key, *args):
  if not config._check_vma.value or not config.auto_pcast.value:
    return key, args
  if not args:
    return key, args
  key_vma = core.typeof(key).mat.varying
  out = []
  for a in args:
    arg_vma = (aval.mat.varying
               if isinstance(aval := core.typeof(a), core.ShapedArray)
               else frozenset())
    # If key is less varying than the args, then it's an error and user should
    # pvary at their level because it has key-reuse implications. They can
    # shard the keys passed to shard_map correctly so as to avoid key-reuse
    # getting correctly varying keys. But JAX shouldn't auto-pvary the key.
    if key_vma - arg_vma:
      a = core.pvary(a, tuple(k for k in key_vma if k not in arg_vma))
    if core.typeof(key).mat != core.typeof(a).mat:
      raise TypeError(
          f"{name} requires all arguments to have matching type. Got key type:"
          f" {core.typeof(key)} vs arg type: {core.typeof(a)}. Use"
          " jax.lax.pcast(..., to='varying') to make them match. If your key is"
          " less varying than arg, watch out for key-reuse problems.")
    out.append(a)
  return key, out

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap the mismatched argument (or key) with jax.lax.pcast(x, to='varying') so both sides have matching varying axes
2. Re-split/re-shard the key with the same sharding as the arguments (e.g. jax.device_put(key, same_named_sharding))
3. Restructure so the key is generated inside the shard_map at matching variance, avoiding key reuse

Example fix

# before
out = jax.random.normal(key, scale.shape, dtype=scale.dtype)  # key less varying than scale under shard_map
# after
out = jax.random.normal(jax.lax.pcast(key, to='varying'), scale.shape, dtype=scale.dtype)

def check_sampler_types(name, key, *args):
    import jax._src.core as core
    for a in args:
        if core.typeof(key).mat != core.typeof(a).mat:
            a = jax.lax.pcast(a, to='varying')
    return key, args

def types_match(key, a) -> bool:
    from jax._src import core
    return core.typeof(key).mat == core.typeof(a).mat

try:
    out = jax.random.normal(key, shape)
except TypeError as e:
    if 'matching type' in str(e):
        key = jax.lax.pcast(key, to='varying')
        out = jax.random.normal(key, shape)
    else:
        raise

When it happens

Trigger: Calling a jax.random sampler under shard_map where the key has a different sharding/varying annotation than a parameter argument (e.g. key replicated on one mesh axis while scale varies along it, in mismatched order), so their core.typeof(...).mat differs after auto-pvary.

Common situations: Multi-device training loops where the PRNG key is created outside shard_map with different NamedSharding than the weights; mixing manually pvary-ed inputs with keys from jax.random.split inside a sharded computation; JAX version upgrades that stopped auto-pvarying keys for key-reuse safety.

Related errors

• Mesh must be provided for shard_map with checkify.
• in_specs passed to shard_map: {s} does not match the specs o
• Mosaic kernels cannot be automatically partitioned. Please w
• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/b49b1cd94bd0d346. Report an issue: GitHub.