jax-ml/jax · error · ValueError

Cannot specify both scratch_shapes and scratch_types. Use sc

Error message

Cannot specify both scratch_shapes and scratch_types. Use scratch_types.

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic_gpu/core.py:349

    deprecations.warn(
        "jax-pallas-mgpu-shapes-types",
        "The out_shape and scratch_shapes arguments to plgpu.kernel are"
        " deprecated. Use out_type and scratch_types instead.",
        stacklevel=2,
    )

  if not isinstance(out_shape, api.NotSpecified):
    if not isinstance(out_type, api.NotSpecified):
      raise ValueError(
          "Cannot specify both out_shape and out_type. Use out_type."
      )
    out_type = out_shape
  elif isinstance(out_type, api.NotSpecified):
    out_type = ()

  if not isinstance(scratch_shapes, api.NotSpecified):
    if not isinstance(scratch_types, api.NotSpecified):
      raise ValueError(
          "Cannot specify both scratch_shapes and scratch_types. Use"
          " scratch_types."
      )
    scratch_types = scratch_shapes
  elif isinstance(scratch_types, api.NotSpecified):
    scratch_types = ()

  if unwrap_out := not isinstance(out_type, (tuple, list)):
    out_type = (out_type,)

  mesh = Mesh(
      grid=grid,
      grid_names=grid_names,
      cluster=cluster,
      cluster_names=cluster_names,
      num_threads=num_threads,
      thread_name=thread_name,
      **mesh_kwargs,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove `scratch_shapes=` and pass the same value as `scratch_types=`
2. For multi-version compatibility, detect the supported kwarg via inspect.signature and pass only one

Example fix

# before
kernel_fn = mgpu.kernel(fn, out_type=..., scratch_shapes=[SMEM((64,64), jnp.float32)], scratch_types=[SMEM((64,64), jnp.float32)], grid=grid)
# after
kernel_fn = mgpu.kernel(fn, out_type=..., scratch_types=[SMEM((64,64), jnp.float32)], grid=grid)

def make_kernel(fn, **kw):
    if 'scratch_types' in kw and 'scratch_shapes' in kw:
        del kw['scratch_shapes']
    return mgpu.kernel(fn, **kw)

When it happens

Trigger: Calling `mgpu.kernel(..., scratch_shapes=[...], scratch_types=[...])` with both keyword arguments set to concrete values.

Common situations: 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.

Related errors

• Cannot specify both out_shape and out_type. Use out_type.
• packed cannot be specified if layout is specified.
• packed, collective and layout arguments are only supported f
• Unsupported dtype: {ref.dtype}
• Only byte-aligned shapes are supported. Got shape: {ref.dtyp

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/9dac44b5310fbf92. Report an issue: GitHub.