{"record":{"id":"d4ca297a8c74b809","repo":"jax-ml/jax","slug":"invalid-spec-spec","errorCode":null,"errorMessage":"Invalid spec: {spec}","messagePattern":"Invalid spec: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":433,"sourceCode":"    _check_specs(SpecErrorType.input, in_specs, axis_names)\n    _check_unreduced(SpecErrorType.input, mesh, axis_names, in_specs)\n  _check_specs(SpecErrorType.out, out_specs, axis_names)\n  _check_unreduced(SpecErrorType.out, mesh, axis_names, out_specs)\n  return mesh, axis_names\n\n\ndef _manual_spec(manual_axes, spec: P, mesh) -> P:\n  out: list[str | tuple[str | None, ...] | None] = []\n  s: str | None | tuple[str, ...]\n  for s in spec.partitions:\n    if s is None:\n      out.append(s)\n    elif isinstance(s, tuple):\n      temp = [p if p in manual_axes else None for p in s]\n      while temp and temp[-1] is None:\n        temp.pop()\n      if None in temp:\n        raise ValueError(f\"Invalid spec: {spec}\")\n      out.append(None if len(temp) == 0 else tuple(temp))\n    else:\n      out.append(s if s in manual_axes else None)\n  _check_unreduced(SpecErrorType.input, mesh, manual_axes, spec)\n  return spec.update(partitions=tuple(out))\n\n\n# Error checking and messages\n\nSpecErrorType = enum.Enum('SpecErrorType', ['input', 'out'])\n\ndef _check_unreduced(error_type, mesh, manual_axes, specs):\n  from jax._src.hijax import HiPspec\n  prefix = 'in' if error_type == SpecErrorType.input else 'out'\n  full_manual = frozenset(mesh.axis_names) == manual_axes\n  specs_flat, _ = tree_flatten(specs)\n  for s in specs_flat:\n    if isinstance(s, HiPspec):","sourceCodeStart":415,"sourceCodeEnd":451,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L415-L451","documentation":"In full-manual shard_map mode, _manual_spec converts a PartitionSpec into a manual-mode spec by keeping only entries in manual_axes. If a tuple spec contains None in a position that isn't trailing (after stripping trailing Nones), the spec is malformed for manual conversion and this ValueError 'Invalid spec: {spec}' is raised.","triggerScenarios":"Using shard_map in manual_axes mode with a PartitionSpec like P('i', None, 'j') where an inner None remains after trailing-None removal while the surrounding entries are manual axes — the conversion can't map the hole to a single manual axis.","commonSituations":"Writing manual-mode sharding code (FSDP/pipeline-style) where specs mix mesh axes and unconstrained dims; porting spec patterns from automatic mode into full manual mode.","solutions":["Rewrite the spec so non-manual entries are only trailing: P('i','j') or P('i','j',None)","Use unreduced/reduced spec syntax or explicit full-manual specs appropriate to manual_axes mode instead of mixing inner Nones"],"exampleFix":"// before\njax.shard_map(f, mesh=mesh, in_specs=P('i', None, 'j'), manual_axes=('i','j'), ...)\n\n// after\njax.shard_map(f, mesh=mesh, in_specs=P('i', 'j'), manual_axes=('i','j'), ...)","handlingStrategy":"validation","validationCode":"def valid_manual_spec(spec, manual_axes):\n    if isinstance(spec, tuple):\n        core = [p for p in spec if p is not None]\n        return all(p in manual_axes for p in core)\n    return True\nassert valid_manual_spec(P('i', None, 'j'), manual_axes), 'inner None not allowed'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["In manual mode keep non-manual dims only as trailing Nones","Test specs against manual_axes before passing to shard_map"],"tags":["jax","shard-map","partition-spec","manual-mode"],"backgroundTag":"invalid-partition-spec","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}