{"record":{"id":"9dac44b5310fbf92","repo":"jax-ml/jax","slug":"cannot-specify-both-scratch-shapes-and-scratch-typ","errorCode":null,"errorMessage":"Cannot specify both scratch_shapes and scratch_types. Use scratch_types.","messagePattern":"Cannot specify both scratch_shapes and scratch_types\\. Use scratch_types\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":349,"sourceCode":"    deprecations.warn(\n        \"jax-pallas-mgpu-shapes-types\",\n        \"The out_shape and scratch_shapes arguments to plgpu.kernel are\"\n        \" deprecated. Use out_type and scratch_types instead.\",\n        stacklevel=2,\n    )\n\n  if not isinstance(out_shape, api.NotSpecified):\n    if not isinstance(out_type, api.NotSpecified):\n      raise ValueError(\n          \"Cannot specify both out_shape and out_type. Use out_type.\"\n      )\n    out_type = out_shape\n  elif isinstance(out_type, api.NotSpecified):\n    out_type = ()\n\n  if not isinstance(scratch_shapes, api.NotSpecified):\n    if not isinstance(scratch_types, api.NotSpecified):\n      raise ValueError(\n          \"Cannot specify both scratch_shapes and scratch_types. Use\"\n          \" scratch_types.\"\n      )\n    scratch_types = scratch_shapes\n  elif isinstance(scratch_types, api.NotSpecified):\n    scratch_types = ()\n\n  if unwrap_out := not isinstance(out_type, (tuple, list)):\n    out_type = (out_type,)\n\n  mesh = Mesh(\n      grid=grid,\n      grid_names=grid_names,\n      cluster=cluster,\n      cluster_names=cluster_names,\n      num_threads=num_threads,\n      thread_name=thread_name,\n      **mesh_kwargs,","sourceCodeStart":331,"sourceCodeEnd":367,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L331-L367","documentation":"In `jax.experimental.pallas.mosaic_gpu.kernel`, `scratch_shapes` is the deprecated alias of `scratch_types`. Passing both at once raises this ValueError; use `scratch_types` only.","triggerScenarios":"Calling `mgpu.kernel(..., scratch_shapes=[...], scratch_types=[...])` with both keyword arguments set to concrete values.","commonSituations":"Upgrading Pallas kernels across JAX versions where scratch_shapes was renamed scratch_types; codemods adding the new kwarg while leaving the old one; copied examples mixing old and new API styles.","solutions":["Remove `scratch_shapes=` and pass the same value as `scratch_types=`","For multi-version compatibility, detect the supported kwarg via inspect.signature and pass only one"],"exampleFix":"# before\nkernel_fn = mgpu.kernel(fn, out_type=..., scratch_shapes=[SMEM((64,64), jnp.float32)], scratch_types=[SMEM((64,64), jnp.float32)], grid=grid)\n# after\nkernel_fn = mgpu.kernel(fn, out_type=..., scratch_types=[SMEM((64,64), jnp.float32)], grid=grid)","handlingStrategy":"validation","validationCode":"def make_kernel(fn, **kw):\n    if 'scratch_types' in kw and 'scratch_shapes' in kw:\n        del kw['scratch_shapes']\n    return mgpu.kernel(fn, **kw)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never mix scratch_shapes and scratch_types","Codemod away legacy kwarg names when bumping JAX versions"],"tags":["jax","pallas","mosaic-gpu","deprecation","scratch-memory"],"backgroundTag":"deprecated-argument-conflict","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}