{"record":{"id":"7e8c5aeafb00a98c","repo":"jax-ml/jax","slug":"pallas-call-requires-all-mesh-axes-to-be-manual-g","errorCode":null,"errorMessage":"pallas_call requires all mesh axes to be Manual, got {get_abstract_mesh().axis_types}","messagePattern":"pallas_call requires all mesh axes to be Manual, got (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/pallas_call.py","lineNumber":135,"sourceCode":"  # Make sure we don't return ShapedArray with pallas memory space to the\n  # outside world.\n  out_avals = tuple(a.update(memory_space=jax_core.MemorySpace.Device)\n                    if isinstance(a, jax_core.ShapedArray) else a\n                    for a in out_avals)\n\n  # TODO(mattjj,yashkatariya): if we hide vmapped away mesh axes, use this:\n  # if not (all(a.sharding.mesh.are_all_axes_manual for a in avals) and\n  #         all(a.sharding.mesh.are_all_axes_manual for a in out_avals) and\n  #         get_abstract_mesh().are_all_axes_manual):\n  #   raise ValueError(\"pallas_call requires all mesh axes to be Manual, \"\n  #                    f\"got {get_abstract_mesh().axis_types}\")\n\n  # NOTE(mattjj,yashkatariya): this doesn't catch auto-mode non-manual axes\n  if not (all(p is None for a in avals if isinstance(a, jax_core.ShapedArray)\n              for p in a.sharding.spec) and\n          all(p is None for a in out_avals if isinstance(a, jax_core.ShapedArray)\n              for p in a.sharding.spec)):\n    raise ValueError(\"pallas_call requires all mesh axes to be Manual, \"\n                     f\"got {get_abstract_mesh().axis_types}\")\n  return out_avals, effs\n\n\npallas_call_p.def_effectful_abstract_eval(_pallas_call_abstract_eval)\n\ndef _pallas_call_is_high(*_, jaxpr, **params):\n  del params\n  return jaxpr.is_high\npallas_call_p.is_high = _pallas_call_is_high\n\n\ndef _get_index_mapping(avals) -> dict[int, tuple[int, ...]]:\n  indices = {}\n  counter = 0\n  for i, in_aval in enumerate(avals):\n    local_counter = []\n    for _ in range(len(in_aval.lo_ty())):","sourceCodeStart":117,"sourceCodeEnd":153,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/pallas_call.py#L117-L153","documentation":"pallas_call operates on manual (per-device, unsharded) data. During abstract evaluation it checks that no ShapedArray input or output has a non-None sharding spec entry (i.e. no named/auto sharding along mesh axes); if any value is sharded across an axis, it raises ValueError showing the abstract mesh's axis types. The note in source acknowledges this check does not catch auto-mode non-manual axes at this point.","triggerScenarios":"Calling pallas_call with arguments or out_shape values that carry a NamedSharding/PartitionSpec placing data on mesh axes, while the abstract mesh has Named or Auto axes (not fully Manual).","commonSituations":"Using pallas_call inside jax.jit with a mesh context (e.g. via mesh() context manager or sharding annotations) intended for sharded SPMD code; forgetting to enter a manual mesh scope before calling Pallas kernels; mixing new sharding APIs with Pallas.","solutions":["Wrap the pallas_call invocation in a manual-replicated/manual mesh scope (e.g. with mesh(manual_axes...): or jax.lax.with_sharding_constraint-free manual context) so all axes are Manual","Remove NamedSharding/PartitionSpec annotations from the arrays and out_shape passed to pallas_call","Run pallas_call outside the sharded/jit-with-mesh context, operating on fully replicated per-device buffers"],"exampleFix":"# before\nwith mesh(mesh_obj, ('x',)):  # Named axis\n    pallas_call(kernel, grid=grid, out_shape=out_shape)(x)\n# after\nwith mesh(mesh_obj):  # all axes manual\n    pallas_call(kernel, grid=grid, out_shape=out_shape)(x)","handlingStrategy":"validation","validationCode":"from jax.sharding import NamedSharding\n\ndef all_manual(vals):\n    for v in vals:\n        s = getattr(v, 'sharding', None)\n        if s is not None and any(p is not None for p in s.spec):\n            return False\n    return True\nassert all_manual(list(args) + list(out_shapes))","typeGuard":"def is_unsharded(x) -> bool:\n    s = getattr(x, 'sharding', None)\n    return s is None or all(p is None for p in s.spec)","tryCatchPattern":null,"preventionTips":["Enter a fully-manual mesh scope before any pallas_call","Keep Pallas kernels operating on replicated per-device buffers; never pass NamedSharding-annotated arrays"],"tags":["jax","pallas","pallas-call","sharding","mesh","manual-axes"],"backgroundTag":"sharded-input-to-manual-kernel","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}