{"record":{"id":"fb855231aba4a4a3","repo":"jax-ml/jax","slug":"semaphore-value-shape-ref-value-shape-does-not-m","errorCode":null,"errorMessage":"Semaphore value shape {ref_value.shape} does not match aval shape {ref_aval.shape}","messagePattern":"Semaphore value shape (.+?) does not match aval shape (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/primitives.py","lineNumber":914,"sourceCode":"  sem_dtype = sem_aval.dtype\n  if not any(\n      jnp.issubdtype(sem_dtype, sem_type)\n      for sem_type in allowed_semaphore_types\n  ):\n    raise ValueError(\n        f\"Must {name} semaphores of the following types:\"\n        f\" {allowed_semaphore_types}. Got {sem_dtype}.\"\n    )\n\n\ndef _transform_semaphore(ref_value, transforms, ref_aval):\n  \"\"\"Helper function for indexing into a semaphore during state_discharge.\"\"\"\n  if ref_value.shape == ref_aval.shape:\n    return state_discharge.transform_array(ref_value, transforms)\n  elif len(ref_value.shape) == 0:\n    return ref_value\n  else:\n    raise ValueError(\n        f\"Semaphore value shape {ref_value.shape} does not match aval shape\"\n        f\" {ref_aval.shape}\"\n    )\n\n\nsemaphore_read_p = jax_core.Primitive(\"semaphore_read\")\nsemaphore_read_p.multiple_results = False\n\n\ndef semaphore_read(sem_or_view) -> jax_typing.Array:\n  \"\"\"Reads the value of a semaphore.\n\n  Args:\n    sem_or_view: A Ref (or view) representing a semaphore.\n\n  Returns:\n    A scalar Array containing the value of the semaphore.\n  \"\"\"","sourceCodeStart":896,"sourceCodeEnd":932,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/primitives.py#L896-L932","documentation":"During state discharge of a semaphore, the runtime value's shape must match the ref aval's shape (or be scalar). A mismatch means the stored semaphore value was shaped differently from what the ref abstraction declares.","triggerScenarios":"Interpret/discharge paths where the semaphore value stored has a shape differing from the Ref's declared shape, e.g. after improper allocation or transform of the semaphore.","commonSituations":"Mostly an internal invariant violation seen in interpret mode with incorrectly shaped semaphores; rare user-facing unless custom ref transforms are used.","solutions":["Ensure semaphores are scalar () Refs with the correct semaphore dtype","Avoid custom ref transforms on semaphores","Update JAX / report a bug with a reproducer if shapes are correct"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"assert ref_value.shape == ref_aval.shape or ref_value.shape == ()","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep semaphore values scalar and consistent with ref avals","Avoid custom transforms on semaphore refs"],"tags":["pallas","semaphore","shape-mismatch","state-discharge","jax"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}