{"record":{"id":"b49b1cd94bd0d346","repo":"jax-ml/jax","slug":"name-requires-all-arguments-to-have-matching-typ","errorCode":null,"errorMessage":"{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.","messagePattern":"(.+?) requires all arguments to have matching type\\. Got key type: (.+?) vs arg type: (.+?)\\. Use jax\\.lax\\.pcast\\(\\.\\.\\., to='varying'\\) to make them match\\. If your key is less varying than arg, watch out for key-reuse problems\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/core.py","lineNumber":3783,"sourceCode":"def random_insert_pvary(name, key, *args):\n  if not config._check_vma.value or not config.auto_pcast.value:\n    return key, args\n  if not args:\n    return key, args\n  key_vma = core.typeof(key).mat.varying\n  out = []\n  for a in args:\n    arg_vma = (aval.mat.varying\n               if isinstance(aval := core.typeof(a), core.ShapedArray)\n               else frozenset())\n    # If key is less varying than the args, then it's an error and user should\n    # pvary at their level because it has key-reuse implications. They can\n    # shard the keys passed to shard_map correctly so as to avoid key-reuse\n    # getting correctly varying keys. But JAX shouldn't auto-pvary the key.\n    if key_vma - arg_vma:\n      a = core.pvary(a, tuple(k for k in key_vma if k not in arg_vma))\n    if core.typeof(key).mat != core.typeof(a).mat:\n      raise TypeError(\n          f\"{name} requires all arguments to have matching type. Got key type:\"\n          f\" {core.typeof(key)} vs arg type: {core.typeof(a)}. Use\"\n          \" jax.lax.pcast(..., to='varying') to make them match. If your key is\"\n          \" less varying than arg, watch out for key-reuse problems.\")\n    out.append(a)\n  return key, out\n","sourceCodeStart":3765,"sourceCodeEnd":3790,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/core.py#L3765-L3790","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Wrap the mismatched argument (or key) with jax.lax.pcast(x, to='varying') so both sides have matching varying axes","Re-split/re-shard the key with the same sharding as the arguments (e.g. jax.device_put(key, same_named_sharding))","Restructure so the key is generated inside the shard_map at matching variance, avoiding key reuse"],"exampleFix":"# before\nout = jax.random.normal(key, scale.shape, dtype=scale.dtype)  # key less varying than scale under shard_map\n# after\nout = jax.random.normal(jax.lax.pcast(key, to='varying'), scale.shape, dtype=scale.dtype)","handlingStrategy":"validation","validationCode":"def check_sampler_types(name, key, *args):\n    import jax._src.core as core\n    for a in args:\n        if core.typeof(key).mat != core.typeof(a).mat:\n            a = jax.lax.pcast(a, to='varying')\n    return key, args","typeGuard":"def types_match(key, a) -> bool:\n    from jax._src import core\n    return core.typeof(key).mat == core.typeof(a).mat","tryCatchPattern":"try:\n    out = jax.random.normal(key, shape)\nexcept TypeError as e:\n    if 'matching type' in str(e):\n        key = jax.lax.pcast(key, to='varying')\n        out = jax.random.normal(key, shape)\n    else:\n        raise","preventionTips":["Shard keys identically to arguments (same NamedSharding/mesh axes)","Generate/split keys inside shard_map at the right variance level","Test multi-device code on a small mesh in CI"],"tags":["jax","sharding","shard-map","pvary","random"],"backgroundTag":"sharding-type-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}